From 3be14f8f6f797bd8ce34165abc94e49f7b2722d5 Mon Sep 17 00:00:00 2001 From: zphinx Date: Thu, 14 May 2026 20:00:38 +0200 Subject: [PATCH] commit all of this --- .gitea/workflows/ci.yml | 3 + AGENTS.md | 28 ++ CHANGELOG.md | 28 +- README.md | 111 +++++ ROADMAP.md | 8 +- docs/ARCHITECTURE.md | 15 +- docs/tai.1 | 284 ++++++++++++ src/tai/chroma_telemetry.py | 5 +- src/tai/cli.py | 903 ++++++++++++++++++++++++++++++++---- src/tai/history_store.py | 372 +++++++++++++++ src/tai/runbook_store.py | 40 +- tai-live-ai-check.md | 3 + tests/test_cli.py | 328 ++++++++++++- tests/test_history_store.py | 115 +++++ tests/test_runbook_store.py | 26 ++ 15 files changed, 2138 insertions(+), 131 deletions(-) create mode 100644 docs/tai.1 create mode 100644 src/tai/history_store.py create mode 100644 tai-live-ai-check.md create mode 100644 tests/test_history_store.py diff --git a/.gitea/workflows/ci.yml b/.gitea/workflows/ci.yml index 1687602..384b59e 100644 --- a/.gitea/workflows/ci.yml +++ b/.gitea/workflows/ci.yml @@ -93,5 +93,8 @@ jobs: - name: Type-check run: .venv/bin/python -m mypy src + - name: Validate man-page sync with CLI + run: .venv/bin/python -m pytest tests/test_cli.py::test_man_page_covers_cli_long_options -v + - name: Test run: .venv/bin/python -m pytest diff --git a/AGENTS.md b/AGENTS.md index 15c5384..067ea64 100644 --- a/AGENTS.md +++ b/AGENTS.md @@ -89,3 +89,31 @@ Before release: - Do not force Kubernetes deployment guidance for current architecture. - Treat Docker as one-shot execution model with mounted persistent volumes for runbooks/sessions/logs. + +## CI Pipeline: Man-Page Validation + +The man-page drift detection is **automatically integrated** into `.gitea/workflows/ci.yml`: + +- **Test name:** `test_man_page_covers_cli_long_options()` in `tests/test_cli.py` +- **Trigger:** Runs on every push and pull request (as part of the `Test` step) +- **Behavior:** Extracts all long options from `tai --help` and verifies each is documented in `docs/tai.1` +- **Failure mode:** CI fails if any long option in CLI is missing from man page; prevents merge of undocumented options + +### How to Fix Failed Man-Page Validation + +1. **Identify missing options:** CI output shows "Missing options in docs/tai.1: ..." +2. **Update docs/tai.1:** Add the missing option to the appropriate section (command or global options) +3. **Re-run tests locally:** `python -m pytest tests/test_cli.py::test_man_page_covers_cli_long_options -v` +4. **Push to trigger CI:** Once local test passes, push the update to trigger CI validation + +### Man-Page Maintenance Workflow + +- **When adding CLI options:** Add to `src/tai/cli.py` and immediately update `docs/tai.1` in the same commit +- **When removing CLI options:** Remove from `src/tai/cli.py` and update `docs/tai.1` in the same commit +- **When renaming CLI options:** Update both CLI code and `docs/tai.1` in one commit +- **When changing option behavior/defaults:** Update the option description in `docs/tai.1` to reflect new behavior + +## Documentation Maintenance + +- When adding or changing CLI commands/options, update `docs/tai.1` in the same change. +- Keep `README.md` and `docs/tai.1` aligned for user-facing flags and examples. diff --git a/CHANGELOG.md b/CHANGELOG.md index 3c96a0d..71d7a6f 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -10,24 +10,38 @@ ______________________________________________________________________ ### Added -- Tier 3 core session memory implementation: - - new `src/tai/session_store.py` persistent ChromaDB store - - `--session-memory` option on `tai run` - - prior-session retrieval injected into analysis/follow-up prompts - - final response indexing at session end +- Unified persistent run history in SQLite: + - new `src/tai/history_store.py` for host-scoped JSON run records + - `--history-db` and `--history/--no-history` options on `tai run` + - prior host history auto-loaded for analysis/follow-up prompts + - analyzed runs auto-indexed into history DB +- External database targets for history and runbook options: + - `--history-db` now supports SQLite path/URL and PostgreSQL DSN + - `--runbooks`/`--store` now support remote ChromaDB URLs +- External DB authentication options: + - history DB: `--history-db-user`, `--history-db-password`, `--env-file` + - runbook store: `--runbooks-user`/`--runbooks-password` and `--store-user`/`--store-password` + - dotenv credential keys: `TAI_HISTORY_DB_USER`, `TAI_HISTORY_DB_PASSWORD`, `TAI_RUNBOOK_STORE_USER`, `TAI_RUNBOOK_STORE_PASSWORD` +- Remote runbook/playbook source ingestion: + - `tai runbooks sync --path` now supports `ssh://` directories + - `tai runbooks sync --path` now supports HTTP/HTTPS webroots with Markdown links + - `tai runbooks add` now supports `ssh://` and HTTP/HTTPS Markdown files - `--output-file` option on `tai run` to persist final AI analysis output as Markdown - `--output-format markdown|json` for `--output-file` exports +- JSON export schema now includes host-specific run metadata (`generated_at`, collection stats, token usage) +- New SQLite run history database (`--history-db`) now stores per-run JSON payloads and auto-loads prior host history for analysis context - Planner enhancements for broader service detection: - generic service candidate extraction from free text - package presence probes in plans (`rpm -q` and `dpkg-query -W`) - SSH read-only allowlist expanded to permit package presence commands (`rpm`, `dpkg-query`) -- Session memory tests in `tests/test_session_store.py` +- History DB tests in `tests/test_history_store.py` - CLI test coverage for analysis output file writing (`tests/test_cli.py`) - CLI test coverage for JSON export and ANSI stripping in written output (`tests/test_cli.py`) ### Changed -- Documentation alignment updates in README and ROADMAP to reflect implemented session memory and package-presence capabilities. +- History reads/writes are now unified on SQLite DB in CLI workflows (`history`, interactive `/history`, analysis context injection). +- Documentation alignment updates in README and ROADMAP to reflect implemented history DB and package-presence capabilities. - Package version metadata alignment: `src/tai/__init__.py` now matches project version `0.4.0`. ______________________________________________________________________ diff --git a/README.md b/README.md index f6f5dc9..c9df89f 100644 --- a/README.md +++ b/README.md @@ -35,6 +35,7 @@ The tool may suggest remediation commands in output, but does not execute them. - Diagnostics collection mode - AI analysis mode - Optional analysis export via `--output-file ` (`--output-format markdown|json`) +- Automatic host history persistence/read via database (`--history-db`, `--history/--no-history`) - Interactive loop with `/collect`, `/analyze`, `/help`, `/quit` ### AI and Prompting @@ -184,6 +185,106 @@ tai run "sshd authentication failed" \ --output-format json ``` +JSON export includes host-specific run metadata: + +- `schema` and `generated_at` +- `issue`, `host`, `model` +- `collection` summary (`total`, `failed`, `succeeded`) +- `token_usage` (`prompt_tokens`, `completion_tokens`, `total_tokens`) when available from backend +- `analysis` text + +By default, each analyzed run is also written to the history database and prior +sessions for the same host are read and injected as historical context. + +Database targets supported by `--history-db`: + +- SQLite file path (for example `~/.tai/history.db`) +- SQLite URL (for example `sqlite:////tmp/tai-history.db`) +- PostgreSQL DSN (for example `postgresql://user:pass@dbhost:5432/tai`) + +Example using remote PostgreSQL history database: + +```bash +tai run "sshd authentication failed" \ + --host bastion01 \ + --collect --analyze \ + --history-db postgresql://tai_user:secret@db.internal:5432/tai +``` + +Credential options for external history DB: + +- `--history-db-user ` +- `--history-db-password ` +- `--env-file ` (loads dotenv values) + +Dotenv keys for history DB credentials: + +- `TAI_HISTORY_DB_USER` +- `TAI_HISTORY_DB_PASSWORD` + +Runbook store targets supported by `--runbooks` and `tai runbooks --store`: + +- Local embedded ChromaDB path (default) +- Remote ChromaDB URL (for example `http://chroma.internal:8000`) + +Example using remote ChromaDB runbook store at analysis time: + +```bash +tai run "nginx failing after reboot" \ + --host web01 \ + --collect --analyze \ + --runbooks http://chroma.internal:8000 +``` + +Credential options for remote runbook store: + +- `--runbooks-user ` / `--runbooks-password ` on `tai run` +- `--store-user ` / `--store-password ` on `tai runbooks ...` +- `--env-file ` (loads dotenv values) + +Dotenv keys for runbook store credentials: + +- `TAI_RUNBOOK_STORE_USER` +- `TAI_RUNBOOK_STORE_PASSWORD` + +Remote runbook (playbook) sources supported by `tai runbooks sync --path`: + +- Local directory path (for example `./runbooks`) +- SSH directory URI (for example `ssh://ops@ssh.archflux.net/opt/tai/runbooks`) +- HTTP/HTTPS webroot URL that exposes `.md` links (for example `https://kb.example/runbooks/`) + +Webroot hardening rules: + +- Only `.md` links are considered for download. +- Downloaded payload must look like real Markdown (HTML wrappers are ignored). +- Non-markdown payloads are discarded. +- Downloaded content is never executed. It is stored as plain text and only parsed for AI retrieval context. + +Single runbook (playbook) sources supported by `tai runbooks add`: + +- Local file path +- SSH file URI (for example `ssh://ops@ssh.archflux.net/opt/tai/runbooks/nginx.md`) +- HTTP/HTTPS URL to a Markdown file + +For HTTP/HTTPS single-file add, the source URL must end in `.md` and resolve to Markdown content. + +Examples: + +```bash +# Sync from SSH-hosted runbooks directory into remote ChromaDB +tai runbooks sync \ + --path ssh://ops@ssh.archflux.net/opt/tai/runbooks \ + --store http://chroma.internal:8000 + +# Sync from HTTPS webroot listing Markdown runbooks +tai runbooks sync \ + --path https://kb.example/runbooks/ \ + --store ~/.tai/runbooks + +# Add one runbook directly from HTTPS +tai runbooks add https://kb.example/runbooks/nginx.md --store ~/.tai/runbooks +``` + ## Runbook Workflow 1. Write Markdown runbooks in `runbooks/` with frontmatter keys: `service`, `symptoms`, `tags`. @@ -209,6 +310,16 @@ Focused suites: pytest tests/test_plan.py tests/test_ai.py tests/test_cli.py ``` +## Man Page + +A manual page is available at `docs/tai.1`. + +Render it locally: + +```bash +man ./docs/tai.1 +``` + ## Known Limits - Deep service-specific probes (known binary/config/package aliases) are richer for recognized services than generic service names. diff --git a/ROADMAP.md b/ROADMAP.md index c5dff0b..1976f4b 100644 --- a/ROADMAP.md +++ b/ROADMAP.md @@ -494,10 +494,10 @@ and retention/persistence policies are stable and proven in production usage. ### Suggested Delivery Phases 1. Build baseline Rust CLI scaffold with feature-flagged parity checkpoints -2. Port SSH execution and read-only policy enforcement modules -3. Port planner, collectors, prompt composition, and AI client adapters -4. Port session memory/history and runbook workflows with migration tests -5. Port interactive UX/TUI layer and deprecate Python runtime path +1. Port SSH execution and read-only policy enforcement modules +1. Port planner, collectors, prompt composition, and AI client adapters +1. Port session memory/history and runbook workflows with migration tests +1. Port interactive UX/TUI layer and deprecate Python runtime path ### Rust Toolchain End-State diff --git a/docs/ARCHITECTURE.md b/docs/ARCHITECTURE.md index fa2be9e..176191f 100644 --- a/docs/ARCHITECTURE.md +++ b/docs/ARCHITECTURE.md @@ -53,9 +53,22 @@ This document describes tai's current runtime architecture, module responsibilit ## Data Stores -- Runbook store (Tier 2): local ChromaDB path, default `~/.tai/runbooks` +- Runbook store (Tier 2): local ChromaDB path or remote ChromaDB HTTP endpoint (`--runbooks`, `runbooks --store`) +- Run history store (Tier 3): SQLite file/URL or PostgreSQL DSN (`--history-db`) - Session logs: optional JSONL file configured by `--log-file` +External DB auth can be provided by CLI options or dotenv file (`--env-file`) and is resolved without executing downloaded runbook content. + +## Runbook Source Ingestion + +`tai runbooks sync --path` and `tai runbooks add` support runbook/playbook source retrieval from: + +- local filesystem paths +- SSH URIs (`ssh://...`) via read-only remote fetch (`find`, `cat`) +- HTTP/HTTPS URLs (single `.md` file or webroot index with `.md` links) + +Remote source content is materialized into temporary local files, embedded, and then indexed into the target ChromaDB store. + ## Retrieval Layers - Tier 1 (implemented): in-memory semantic retrieval over diagnostic chunks diff --git a/docs/tai.1 b/docs/tai.1 new file mode 100644 index 0000000..c7df080 --- /dev/null +++ b/docs/tai.1 @@ -0,0 +1,284 @@ +.TH TAI 1 "2026-05-11" "tai 0.4.0" "User Commands" +.SH NAME +tai \- read-only Linux troubleshooting assistant with SSH diagnostics and AI analysis +.SH SYNOPSIS +.B tai +.RI [ GLOBAL_OPTIONS ] +.B run +.I ISSUE +.RI [ RUN_OPTIONS ] +.PP +.B tai +.B history +.RI [ HISTORY_OPTIONS ] +.PP +.B tai +.B runbooks +.B sync +.RI [ SYNC_OPTIONS ] +.PP +.B tai +.B runbooks +.B list +.RI [ LIST_OPTIONS ] +.PP +.B tai +.B runbooks +.B add +.I FILE +.RI [ ADD_OPTIONS ] +.SH DESCRIPTION +.B tai +connects to Linux hosts over SSH, collects read-only diagnostics, and can ask an OpenAI-compatible model for grounded analysis. +.PP +Remote runbook (playbook) sources can be local paths, SSH URIs, or HTTP/HTTPS webroots. +Downloaded runbook content is never executed. It is stored as text and parsed for retrieval context only. +.SH COMMANDS +.SS run +Main troubleshooting entrypoint. +.TP +.BI --host " HOST" +Target host to troubleshoot. +.TP +.BI --port " PORT" +SSH port (default: 22). +.TP +.BI --path " PATH" +Target path to inspect. Repeatable. +.TP +.BI --identity-file " FILE" +SSH private key path. +.TP +.BI --jump-host " HOST" +SSH bastion/jump host. +.TP +.B --ignore-ssh-config +Ignore ~/.ssh/config and rely only on CLI options. +.TP +.B --probe / --no-probe +Enable or disable connectivity probe. +.TP +.B --collect / --no-collect +Collect baseline diagnostics. +.TP +.B --analyze / --no-analyze +Send diagnostics to AI for analysis. +.TP +.B --interactive / --no-interactive +Interactive follow-up mode. +.TP +.BI --ai-host " URL" +OpenAI-compatible AI backend URL. +.TP +.BI --model " NAME" +Model name for analysis. +.TP +.BI --ai-key " KEY" +API key for AI backend. +.TP +.BI --ai-timeout-seconds " SECONDS" +Timeout for AI requests. +.TP +.BI --ai-max-tokens " TOKENS" +Max completion tokens. +.TP +.BI --embed-model " NAME" +Embedding model for RAG. +.TP +.B --no-rag +Disable RAG retrieval. +.TP +.B --rag-debug / --no-rag-debug +Print retrieval debug output. +.TP +.BI --runbooks " STORE" +Runbook store path or remote Chroma URL. +.TP +.BI --runbooks-user " USER" +Runbook store login/user for remote Chroma URLs. +.TP +.BI --runbooks-password " PASSWORD" +Runbook store password for remote Chroma URLs. +.TP +.BI --history-db " TARGET" +History DB target: SQLite path/URL or PostgreSQL DSN. +.TP +.BI --history-db-user " USER" +History DB login/user for external database URLs. +.TP +.BI --history-db-password " PASSWORD" +History DB password for external database URLs. +.TP +.B --history / --no-history +Enable or disable history DB reads/writes. +.TP +.BI --output-file " FILE" +Write analysis to file. +.TP +.BI --output-format " FORMAT" +Output format: markdown or json. +.TP +.BI --log-file " FILE" +Optional JSONL event log path. +.TP +.BI --env-file " FILE" +Optional dotenv file used to resolve DB credentials. +.SS history +Search/list indexed troubleshooting history. +.TP +.BI --query " TEXT" +Optional keyword search in issue/summary. +.TP +.BI --host " HOST" +Filter by host. +.TP +.BI --limit " N" +Maximum sessions to show. +.TP +.BI --export " FILE" +Export results as Markdown. +.TP +.BI --history-db " TARGET" +History DB target: SQLite path/URL or PostgreSQL DSN. +.TP +.BI --history-db-user " USER" +History DB login/user for external database URLs. +.TP +.BI --history-db-password " PASSWORD" +History DB password for external database URLs. +.TP +.BI --env-file " FILE" +Optional dotenv file used to resolve DB credentials. +.SS runbooks sync +Index all runbooks from source path. +.TP +.BI --path " SOURCE" +Runbook source: local directory, ssh://host/path, or http(s) webroot URL. +.TP +.BI --store " TARGET" +ChromaDB store path or remote URL. +.TP +.BI --store-user " USER" +Runbook store login/user for remote Chroma URLs. +.TP +.BI --store-password " PASSWORD" +Runbook store password for remote Chroma URLs. +.TP +.BI --ai-host " URL" +OpenAI-compatible AI backend URL. +.TP +.BI --embed-model " NAME" +Embedding model name. +.TP +.BI --ai-key " KEY" +API key for AI backend. +.TP +.BI --identity-file " FILE" +SSH private key for ssh:// source. +.TP +.BI --jump-host " HOST" +SSH bastion for ssh:// source. +.TP +.B --ignore-ssh-config +Ignore ~/.ssh/config for ssh:// source. +.TP +.BI --env-file " FILE" +Optional dotenv file used to resolve DB credentials. +.SS runbooks list +List indexed runbooks. +.TP +.BI --store " TARGET" +ChromaDB store path or remote URL. +.TP +.BI --store-user " USER" +Runbook store login/user for remote Chroma URLs. +.TP +.BI --store-password " PASSWORD" +Runbook store password for remote Chroma URLs. +.TP +.BI --env-file " FILE" +Optional dotenv file used to resolve DB credentials. +.SS runbooks add +Index one runbook file. +.TP +.BI FILE +Runbook source: local file, ssh://host/path/file.md, or HTTP/HTTPS URL ending in .md. +.TP +.BI --store " TARGET" +ChromaDB store path or remote URL. +.TP +.BI --store-user " USER" +Runbook store login/user for remote Chroma URLs. +.TP +.BI --store-password " PASSWORD" +Runbook store password for remote Chroma URLs. +.TP +.BI --ai-host " URL" +OpenAI-compatible AI backend URL. +.TP +.BI --embed-model " NAME" +Embedding model name. +.TP +.BI --ai-key " KEY" +API key for AI backend. +.TP +.BI --identity-file " FILE" +SSH private key for ssh:// source. +.TP +.BI --jump-host " HOST" +SSH bastion for ssh:// source. +.TP +.B --ignore-ssh-config +Ignore ~/.ssh/config for ssh:// source. +.TP +.BI --env-file " FILE" +Optional dotenv file used to resolve DB credentials. +.SH ENVIRONMENT +The following variables are recognized for DB credentials: +.TP +.B TAI_HISTORY_DB_USER +History DB user when --history-db points to external database. +.TP +.B TAI_HISTORY_DB_PASSWORD +History DB password when --history-db points to external database. +.TP +.B TAI_RUNBOOK_STORE_USER +Runbook store user for remote ChromaDB. +.TP +.B TAI_RUNBOOK_STORE_PASSWORD +Runbook store password for remote ChromaDB. +.SH SECURITY NOTES +.TP +\(bu +SSH diagnostics are validated against read-only command policy. +.TP +\(bu +Web/SSH runbook content is never executed. +.TP +\(bu +Webroot ingestion accepts only Markdown-like payloads and skips HTML/non-markdown wrappers. +.SH FILES +.TP +.I ~/.tai/history.db +Default local history database. +.TP +.I ~/.tai/runbooks +Default local runbook store path. +.SH EXAMPLES +.PP +Analyze with PostgreSQL history DB credentials loaded from .env: +.PP +.nf +$ tai run "sshd auth failed" --host bastion01 --collect --analyze \ + --history-db postgresql://db.internal:5432/tai --env-file ./.env +.fi +.PP +Sync runbooks from HTTPS webroot to remote ChromaDB: +.PP +.nf +$ tai runbooks sync --path https://kb.example/runbooks/ \ + --store https://chroma.internal:8443 --env-file ./.env +.fi +.SH SEE ALSO +.BR README.md , +.BR docs/ARCHITECTURE.md diff --git a/src/tai/chroma_telemetry.py b/src/tai/chroma_telemetry.py index 609e686..0329bb2 100644 --- a/src/tai/chroma_telemetry.py +++ b/src/tai/chroma_telemetry.py @@ -7,16 +7,15 @@ disabled, so tai wires ChromaDB to this no-op client instead. from __future__ import annotations -from typing import override - from chromadb.config import System from chromadb.telemetry.product import ProductTelemetryClient, ProductTelemetryEvent +from typing_extensions import override class NoOpProductTelemetryClient(ProductTelemetryClient): """Telemetry client that intentionally drops all events.""" - def __init__(self, system: System): + def __init__(self, system: System) -> None: super().__init__(system) @override diff --git a/src/tai/cli.py b/src/tai/cli.py index d88a54d..0775e12 100644 --- a/src/tai/cli.py +++ b/src/tai/cli.py @@ -4,10 +4,19 @@ from __future__ import annotations import asyncio import json +import os import re +import shlex +import subprocess import sys +import tempfile +from dataclasses import dataclass +from datetime import UTC, datetime +from pathlib import Path from time import perf_counter from typing import Annotated +from urllib.parse import quote, unquote, urljoin, urlparse +from urllib.request import urlopen import typer from rich.console import Console @@ -19,6 +28,7 @@ from rich.text import Text from tai.ai_client import DEFAULT_AI_HOST, DEFAULT_EMBED_MODEL, DEFAULT_MODEL, AIClient, AIConfig from tai.ai_guardrails import validate_ai_response from tai.collectors import CollectionReport, collect_from_plan +from tai.history_store import DEFAULT_HISTORY_DB_PATH, RunHistoryStore from tai.input_parser import InputValidationError, build_request from tai.models import TroubleshootRequest from tai.plan import plan_from_request @@ -32,7 +42,7 @@ from tai.prompt_builder import ( from tai.rag_retriever import EmbeddedChunk, chunk_report, retrieve_scored from tai.runbook_store import RunbookChunk, RunbookStore from tai.session_log import SessionLogger -from tai.session_store import PastSession, SessionStore +from tai.session_store import PastSession from tai.ssh_client import SSHClient, SSHCommandResult, SSHConnectionConfig, SSHSession app = typer.Typer(no_args_is_help=True, add_completion=False) @@ -41,6 +51,111 @@ app.add_typer(runbooks_app, name="runbooks") console = Console() +@dataclass(slots=True) +class RunResult: + """Outcome details captured from a run for optional file export.""" + + final_response: str | None = None + collection_total: int | None = None + collection_failed: int | None = None + ai_prompt_tokens: int | None = None + ai_completion_tokens: int | None = None + ai_total_tokens: int | None = None + generated_at: str | None = None + + +def _load_env_file(path: str | None) -> dict[str, str]: + """Parse a simple KEY=VALUE dotenv file; missing file yields empty mapping.""" + if not path: + return {} + env_path = Path(path).expanduser().resolve() + if not env_path.is_file(): + return {} + + values: dict[str, str] = {} + for raw in env_path.read_text(encoding="utf-8", errors="replace").splitlines(): + line = raw.strip() + if not line or line.startswith("#"): + continue + if line.startswith("export "): + line = line[7:].strip() + if "=" not in line: + continue + key, value = line.split("=", 1) + key = key.strip() + value = value.strip() + if (value.startswith('"') and value.endswith('"')) or ( + value.startswith("'") and value.endswith("'") + ): + value = value[1:-1] + if key: + values[key] = value + return values + + +def _resolve_secret( + explicit: str | None, env_key: str, env_file_values: dict[str, str] +) -> str | None: + if explicit is not None: + return explicit + if env_key in env_file_values and env_file_values[env_key] != "": + return env_file_values[env_key] + env_value = os.getenv(env_key) + if env_value is not None and env_value != "": + return env_value + return None + + +def _inject_url_credentials( + target: str, + *, + user: str | None, + password: str | None, + schemes: set[str], +) -> str: + """Inject username/password into URL netloc for supported schemes.""" + parsed = urlparse(target) + if parsed.scheme.lower() not in schemes: + return target + if user is None and password is None: + return target + + effective_user = user if user is not None else parsed.username + effective_pass = password if password is not None else parsed.password + if effective_user is None and effective_pass is None: + return target + + host = parsed.hostname or "" + if not host: + return target + + user_info = "" + if effective_user is not None: + user_info += quote(effective_user, safe="") + if effective_pass is not None: + user_info += f":{quote(effective_pass, safe='')}" + if user_info: + user_info += "@" + + netloc = f"{user_info}{host}" + if parsed.port is not None: + netloc += f":{parsed.port}" + + return parsed._replace(netloc=netloc).geturl() + + +def _redact_url_password(target: str) -> str: + parsed = urlparse(target) + if parsed.password is None: + return target + return _inject_url_credentials( + target, + user=parsed.username, + password="****", + schemes={parsed.scheme.lower()}, + ) + + @app.command("history") def history( query: Annotated[ @@ -59,21 +174,54 @@ def history( str | None, typer.Option("--export", help="Optional path to write results as Markdown."), ] = None, - session_memory_path: Annotated[ + history_db: Annotated[ str, typer.Option( - "--session-memory", - help="Path to persistent session memory store. Defaults to ~/.tai/sessions.", + "--history-db", + help="History DB target: local SQLite path/URL or PostgreSQL DSN.", ), - ] = "~/.tai/sessions", + ] = DEFAULT_HISTORY_DB_PATH, + history_db_user: Annotated[ + str | None, + typer.Option("--history-db-user", help="History DB login/user for external database URLs."), + ] = None, + history_db_password: Annotated[ + str | None, + typer.Option( + "--history-db-password", + help="History DB password for external database URLs.", + ), + ] = None, + env_file: Annotated[ + str | None, + typer.Option("--env-file", help="Optional dotenv file to resolve DB credentials."), + ] = None, ) -> None: """Search or list previously indexed troubleshooting sessions.""" from pathlib import Path + env_values = _load_env_file(env_file) + resolved_history_user = _resolve_secret( + history_db_user, + "TAI_HISTORY_DB_USER", + env_values, + ) + resolved_history_password = _resolve_secret( + history_db_password, + "TAI_HISTORY_DB_PASSWORD", + env_values, + ) + resolved_history_db = _inject_url_credentials( + history_db, + user=resolved_history_user, + password=resolved_history_password, + schemes={"postgres", "postgresql"}, + ) + try: - store = SessionStore(session_memory_path) + store = RunHistoryStore(resolved_history_db) except Exception as exc: # noqa: BLE001 - console.print(f"[red]Could not open session memory:[/red] {exc}") + console.print(f"[red]Could not open history DB:[/red] {exc}") raise typer.Exit(code=1) from exc try: @@ -218,6 +366,31 @@ def run( case_sensitive=False, ), ] = "markdown", + history_db: Annotated[ + str, + typer.Option( + "--history-db", + help="History DB target: local SQLite path/URL or PostgreSQL DSN.", + ), + ] = DEFAULT_HISTORY_DB_PATH, + history_db_user: Annotated[ + str | None, + typer.Option("--history-db-user", help="History DB login/user for external database URLs."), + ] = None, + history_db_password: Annotated[ + str | None, + typer.Option( + "--history-db-password", + help="History DB password for external database URLs.", + ), + ] = None, + use_history: Annotated[ + bool, + typer.Option( + "--history/--no-history", + help="Enable or disable automatic host history read/write via history DB.", + ), + ] = True, no_rag: Annotated[ bool, typer.Option( @@ -246,16 +419,21 @@ def run( help="Path to a synced runbook ChromaDB store. Enables Tier 2 RAG.", ), ] = None, - session_memory_path: Annotated[ + runbooks_user: Annotated[ + str | None, + typer.Option("--runbooks-user", help="Runbook store login/user for remote ChromaDB URLs."), + ] = None, + runbooks_password: Annotated[ str | None, typer.Option( - "--session-memory", - help=( - "Path to persistent session memory store for prior-session retrieval " - "(Tier 4). Omit to disable." - ), + "--runbooks-password", + help="Runbook store password for remote ChromaDB URLs.", ), ] = None, + env_file: Annotated[ + str | None, + typer.Option("--env-file", help="Optional dotenv file to resolve DB credentials."), + ] = None, ) -> None: """Start an interactive troubleshooting session scaffold.""" try: @@ -303,28 +481,65 @@ def run( if analyze or interactive: console.print(f"[cyan]AI:[/cyan] {AIClient(ai_config).summary()}") + env_values = _load_env_file(env_file) + resolved_history_user = _resolve_secret( + history_db_user, + "TAI_HISTORY_DB_USER", + env_values, + ) + resolved_history_password = _resolve_secret( + history_db_password, + "TAI_HISTORY_DB_PASSWORD", + env_values, + ) + resolved_runbooks_user = _resolve_secret( + runbooks_user, + "TAI_RUNBOOK_STORE_USER", + env_values, + ) + resolved_runbooks_password = _resolve_secret( + runbooks_password, + "TAI_RUNBOOK_STORE_PASSWORD", + env_values, + ) + + resolved_history_db = _inject_url_credentials( + history_db, + user=resolved_history_user, + password=resolved_history_password, + schemes={"postgres", "postgresql"}, + ) + display_history_db = _redact_url_password(resolved_history_db) + display_runbooks_path = _redact_url_password(runbooks_path) if runbooks_path else None + runbook_store: RunbookStore | None = None if runbooks_path is not None: try: - runbook_store = RunbookStore(runbooks_path) + runbook_store = RunbookStore( + runbooks_path, + username=resolved_runbooks_user, + password=resolved_runbooks_password, + ) rb_count = runbook_store.count() - console.print(f"[dim]Runbooks: {rb_count} indexed at {runbooks_path}[/dim]") + console.print(f"[dim]Runbooks: {rb_count} indexed at {display_runbooks_path}[/dim]") except Exception as exc: # noqa: BLE001 console.print(f"[yellow]Runbook store unavailable:[/yellow] {exc}") - session_store: SessionStore | None = None - if session_memory_path: + history_store: RunHistoryStore | None = None + if use_history: try: - session_store = SessionStore(session_memory_path) - mem_count = session_store.count() - console.print( - f"[dim]Session memory: {mem_count} indexed at {session_memory_path}[/dim]" + history_store = RunHistoryStore(resolved_history_db) + host_count = history_store.count(host=req.host) + msg = ( + f"[dim]History DB: {host_count} prior run(s) " + f"for host={req.host} at {display_history_db}[/dim]" ) + console.print(msg) except Exception as exc: # noqa: BLE001 - console.print(f"[yellow]Session memory unavailable:[/yellow] {exc}") + console.print(f"[yellow]History DB unavailable:[/yellow] {exc}") try: - final_response = asyncio.run( + run_result = asyncio.run( _async_main( config, req, @@ -336,21 +551,27 @@ def run( no_rag=no_rag, rag_debug=rag_debug, runbook_store=runbook_store, - session_store=session_store, + history_store=history_store, logger=logger, ) ) if output_file is not None: - if final_response is None: + if run_result.final_response is None: console.print("[yellow]No AI analysis output available to write.[/yellow]") else: _write_analysis_output( output_file, - final_response, + run_result.final_response, output_format=output_format, issue=req.issue, host=req.host, model=model, + collection_total=run_result.collection_total, + collection_failed=run_result.collection_failed, + ai_prompt_tokens=run_result.ai_prompt_tokens, + ai_completion_tokens=run_result.ai_completion_tokens, + ai_total_tokens=run_result.ai_total_tokens, + generated_at=run_result.generated_at, ) except typer.Exit: raise @@ -374,9 +595,9 @@ async def _async_main( no_rag: bool, rag_debug: bool, runbook_store: RunbookStore | None, - session_store: SessionStore | None, + history_store: RunHistoryStore | None, logger: SessionLogger | None, -) -> str | None: +) -> RunResult: """Open a single SSH session and run probe / collection / analysis through it.""" client = SSHClient(config) if logger is not None: @@ -393,6 +614,7 @@ async def _async_main( }, ) async with client.connect() as session: + run_result = RunResult() if probe: result = await session.probe() _handle_probe_result(result) @@ -412,6 +634,8 @@ async def _async_main( console.print(f"[cyan]Collecting diagnostics:[/cyan] {len(plan)} commands") report = await collect_from_plan(session, plan) _handle_collection_report(report) + run_result.collection_total = report.total + run_result.collection_failed = report.failed if logger is not None: logger.log_event( "collection_summary", @@ -423,6 +647,8 @@ async def _async_main( initial_response: str | None = None if analyze and report is not None: + initial_tokens: dict[str, int] = {} + initial_host_history = _query_host_history(history_store, req.host, limit=5) initial_response = _run_analysis( ai_config, req.issue, @@ -430,13 +656,19 @@ async def _async_main( no_rag=no_rag, rag_debug=rag_debug, runbook_store=runbook_store, - session_store=session_store, + host_history=initial_host_history, logger=logger, + token_metrics_sink=initial_tokens, ) + if initial_tokens: + run_result.ai_prompt_tokens = initial_tokens.get("prompt_tokens") + run_result.ai_completion_tokens = initial_tokens.get("completion_tokens") + run_result.ai_total_tokens = initial_tokens.get("total_tokens") interactive_response: str | None = None + interactive_tokens: dict[str, int] | None = None if interactive: - interactive_response = await _interactive_loop( + interactive_response, interactive_tokens = await _interactive_loop( session, req, ai_config, @@ -444,14 +676,32 @@ async def _async_main( no_rag=no_rag, rag_debug=rag_debug, runbook_store=runbook_store, - session_store=session_store, + history_store=history_store, logger=logger, ) + if interactive_tokens: + run_result.ai_prompt_tokens = interactive_tokens.get("prompt_tokens") + run_result.ai_completion_tokens = interactive_tokens.get("completion_tokens") + run_result.ai_total_tokens = interactive_tokens.get("total_tokens") final_response = interactive_response or initial_response - if session_store is not None and final_response: - _index_session_memory(session_store, ai_config, req, final_response, logger=logger) - return final_response + if final_response: + payload = _build_analysis_payload( + issue=req.issue, + host=req.host, + model=ai_config.model, + analysis=final_response, + collection_total=run_result.collection_total, + collection_failed=run_result.collection_failed, + ai_prompt_tokens=run_result.ai_prompt_tokens, + ai_completion_tokens=run_result.ai_completion_tokens, + ai_total_tokens=run_result.ai_total_tokens, + generated_at=run_result.generated_at, + ) + run_result.generated_at = str(payload.get("generated_at", "")) + _index_history_payload(history_store, payload, logger=logger) + run_result.final_response = final_response + return run_result async def _interactive_loop( @@ -463,9 +713,9 @@ async def _interactive_loop( no_rag: bool = False, rag_debug: bool = False, runbook_store: RunbookStore | None = None, - session_store: SessionStore | None = None, + history_store: RunHistoryStore | None = None, logger: SessionLogger | None, -) -> str | None: +) -> tuple[str | None, dict[str, int] | None]: """Run a follow-up loop for collecting and conversational analysis.""" console.print( Panel( @@ -481,6 +731,7 @@ async def _interactive_loop( embedded_chunks: list[EmbeddedChunk] | None = None ai_embed = AIClient(ai_config) last_response: str | None = None + last_tokens: dict[str, int] | None = None if not no_rag and report is not None: embedded_chunks, index_error, index_ms = await asyncio.to_thread( @@ -518,14 +769,14 @@ async def _interactive_loop( else: line = sys.stdin.readline() # non-TTY / piped mode if not line: - return last_response + return last_response, last_tokens command = line.strip() console.print(f"\n[bold cyan]tai[/bold cyan][dim] >[/dim] {command}") except (EOFError, KeyboardInterrupt): console.print("\n[yellow]Exiting interactive mode.[/yellow]") if logger is not None: logger.log_event("interactive_exit", {"reason": "signal_or_eof"}) - return last_response + return last_response, last_tokens if not command: continue @@ -534,7 +785,7 @@ async def _interactive_loop( console.print("[green]Bye.[/green]") if logger is not None: logger.log_event("interactive_exit", {"reason": "user_quit"}) - return last_response + return last_response, last_tokens if command == "/help": console.print( @@ -553,19 +804,19 @@ async def _interactive_loop( continue if command.startswith("/history"): - if session_store is None: + if history_store is None: console.print( - "[yellow]Session memory is disabled. " - "Use --session-memory to enable /history.[/yellow]" + "[yellow]History DB is disabled. " + "Use --history to enable /history.[/yellow]" ) continue keyword = command.removeprefix("/history").strip() try: sessions = ( - session_store.search_keyword(keyword, host=req.host, limit=5) + history_store.search_keyword(keyword, host=req.host, limit=5) if keyword - else session_store.list_recent(host=req.host, limit=5) + else history_store.list_recent(host=req.host, limit=5) ) except Exception as exc: # noqa: BLE001 console.print(f"[yellow]History unavailable:[/yellow] {exc}") @@ -654,6 +905,8 @@ async def _interactive_loop( console.print("[red]No diagnostics available to analyze.[/red]") continue + analyze_tokens: dict[str, int] = {} + analyze_host_history = _query_host_history(history_store, req.host, limit=5) response = _run_followup_analysis( ai_config, req.issue, @@ -663,13 +916,16 @@ async def _interactive_loop( embedded_chunks=embedded_chunks, rag_debug=rag_debug, runbook_store=runbook_store, - session_store=session_store, + host_history=analyze_host_history, logger=logger, + token_metrics_sink=analyze_tokens, ) prior_questions.append("/analyze") if logger is not None: logger.log_event("interactive_followup", {"question": "/analyze"}) last_response = response + if analyze_tokens: + last_tokens = analyze_tokens continue if report is None: @@ -713,6 +969,8 @@ async def _interactive_loop( console.print("[red]No diagnostics available to analyze.[/red]") continue + followup_tokens: dict[str, int] = {} + followup_host_history = _query_host_history(history_store, req.host, limit=5) response = _run_followup_analysis( ai_config, req.issue, @@ -722,13 +980,16 @@ async def _interactive_loop( embedded_chunks=embedded_chunks, rag_debug=rag_debug, runbook_store=runbook_store, - session_store=session_store, + host_history=followup_host_history, logger=logger, + token_metrics_sink=followup_tokens, ) prior_questions.append(command) if logger is not None: logger.log_event("interactive_followup", {"question": command}) last_response = response + if followup_tokens: + last_tokens = followup_tokens def _try_embed_report( @@ -789,8 +1050,9 @@ def _run_analysis( no_rag: bool = False, rag_debug: bool = False, runbook_store: RunbookStore | None = None, - session_store: SessionStore | None = None, + host_history: list[PastSession] | None = None, logger: SessionLogger | None, + token_metrics_sink: dict[str, int] | None = None, ) -> str: """Send collected data to the AI and stream the analysis to stdout.""" console.print() @@ -799,7 +1061,7 @@ def _run_analysis( ai = AIClient(ai_config) system_prompt = build_system_prompt() runbook_chunks = _query_runbooks(runbook_store, issue, ai, top_k=1) - past_sessions = _query_sessions(session_store, issue, report.host, ai, top_k=2) + past_sessions = host_history or [] user_message: str if no_rag: @@ -850,11 +1112,14 @@ def _run_analysis( past_sessions=past_sessions or None, ) try: - response = _complete_ai_response( + response, token_metrics = _complete_ai_response( ai, system_prompt, user_message, ) + if token_metrics_sink is not None and token_metrics is not None: + token_metrics_sink.clear() + token_metrics_sink.update(token_metrics) console.print(Markdown(response)) warnings = validate_ai_response(response) @@ -900,8 +1165,9 @@ def _run_followup_analysis( embedded_chunks: list[EmbeddedChunk] | None = None, rag_debug: bool = False, runbook_store: RunbookStore | None = None, - session_store: SessionStore | None = None, + host_history: list[PastSession] | None = None, logger: SessionLogger | None, + token_metrics_sink: dict[str, int] | None = None, ) -> str: """Run grounded follow-up analysis re-anchored to current diagnostics. @@ -915,7 +1181,7 @@ def _run_followup_analysis( ai = AIClient(ai_config) system_prompt = build_system_prompt() runbook_chunks = _query_runbooks(runbook_store, question, ai, top_k=1) - past_sessions = _query_sessions(session_store, question, report.host, ai, top_k=2) + past_sessions = host_history or [] user_message: str retrieved_names: list[str] = [] @@ -982,11 +1248,14 @@ def _run_followup_analysis( ) try: - response = _complete_ai_response( + response, token_metrics = _complete_ai_response( ai, system_prompt, user_message, ) + if token_metrics_sink is not None and token_metrics is not None: + token_metrics_sink.clear() + token_metrics_sink.update(token_metrics) console.print(Markdown(response)) console.print(Rule(style="dim")) @@ -1018,13 +1287,23 @@ def _complete_ai_response( ai: AIClient, system_prompt: str, user_message: str, -) -> str: +) -> tuple[str, dict[str, int] | None]: """Return a full AI completion in one request. Some local backends intermittently stall on streaming before yielding a first token; using a non-streaming completion path is more reliable for CLI runs. """ - return ai.complete(system_prompt, user_message).content + result = ai.complete(system_prompt, user_message) + prompt_tokens = getattr(result, "prompt_tokens", None) + completion_tokens = getattr(result, "completion_tokens", None) + token_metrics: dict[str, int] | None = None + if isinstance(prompt_tokens, int) and isinstance(completion_tokens, int): + token_metrics = { + "prompt_tokens": prompt_tokens, + "completion_tokens": completion_tokens, + "total_tokens": prompt_tokens + completion_tokens, + } + return result.content, token_metrics def _query_runbooks( @@ -1043,42 +1322,21 @@ def _query_runbooks( return [] -def _query_sessions( - store: SessionStore | None, - question: str, +def _query_host_history( + store: RunHistoryStore | None, host: str, - ai: AIClient, *, - top_k: int = 2, + limit: int, ) -> list[PastSession]: - """Query the session memory store silently; returns empty list on failures.""" + """Return recent host history from the persistent DB; empty on failures.""" if store is None: return [] try: - return store.query(question, host, ai, top_k=top_k) + return store.list_recent(host=host, limit=limit) except Exception: # noqa: BLE001 return [] -def _index_session_memory( - store: SessionStore, - ai_config: AIConfig, - req: TroubleshootRequest, - summary: str, - *, - logger: SessionLogger | None, -) -> None: - """Persist final session summary for future retrieval; non-fatal on failure.""" - try: - ai = AIClient(ai_config) - session_id = store.index_session(req.host, req.issue, summary, ai) - if logger is not None: - logger.log_event("session_memory_indexed", {"session_id": session_id}) - except Exception as exc: # noqa: BLE001 - if logger is not None: - logger.log_event("session_memory_error", {"error": str(exc)}) - - def _format_history_markdown( sessions: list[PastSession], *, @@ -1111,6 +1369,63 @@ def _strip_ansi(text: str) -> str: return ansi_pattern.sub("", text) +def _build_analysis_payload( + *, + issue: str, + host: str, + model: str, + analysis: str, + collection_total: int | None, + collection_failed: int | None, + ai_prompt_tokens: int | None, + ai_completion_tokens: int | None, + ai_total_tokens: int | None, + generated_at: str | None, +) -> dict[str, object]: + """Build a structured JSON payload for export and DB storage.""" + clean_content = _strip_ansi(analysis).rstrip("\n") + succeeded: int | None = None + if collection_total is not None and collection_failed is not None: + succeeded = collection_total - collection_failed + + return { + "schema": "tai.analysis.v1", + "generated_at": generated_at or datetime.now(UTC).isoformat(), + "issue": issue, + "host": host, + "model": model, + "collection": { + "total": collection_total, + "failed": collection_failed, + "succeeded": succeeded, + }, + "token_usage": { + "prompt_tokens": ai_prompt_tokens, + "completion_tokens": ai_completion_tokens, + "total_tokens": ai_total_tokens, + }, + "analysis": clean_content, + } + + +def _index_history_payload( + store: RunHistoryStore | None, + payload: dict[str, object], + *, + logger: SessionLogger | None, +) -> None: + """Store run payload into history DB; non-fatal on failure.""" + if store is None: + return + try: + record_id = store.add_payload(payload) + if logger is not None: + logger.log_event("history_db_indexed", {"record_id": record_id}) + except Exception as exc: # noqa: BLE001 + if logger is not None: + logger.log_event("history_db_error", {"error": str(exc)}) + + def _write_analysis_output( file_path: str, content: str, @@ -1119,6 +1434,12 @@ def _write_analysis_output( issue: str, host: str, model: str, + collection_total: int | None, + collection_failed: int | None, + ai_prompt_tokens: int | None, + ai_completion_tokens: int | None, + ai_total_tokens: int | None, + generated_at: str | None, ) -> None: """Persist final AI analysis output to a file in markdown or json format.""" from pathlib import Path @@ -1128,18 +1449,24 @@ def _write_analysis_output( console.print(f"[red]Invalid --output-format:[/red] {output_format}") raise typer.Exit(code=2) - clean_content = _strip_ansi(content).rstrip() + "\n" + payload = _build_analysis_payload( + issue=issue, + host=host, + model=model, + analysis=content, + collection_total=collection_total, + collection_failed=collection_failed, + ai_prompt_tokens=ai_prompt_tokens, + ai_completion_tokens=ai_completion_tokens, + ai_total_tokens=ai_total_tokens, + generated_at=generated_at, + ) + clean_content = str(payload["analysis"]).rstrip() + "\n" path = Path(file_path).expanduser().resolve() path.parent.mkdir(parents=True, exist_ok=True) if fmt == "json": - payload = { - "issue": issue, - "host": host, - "model": model, - "analysis": clean_content.rstrip("\n"), - } path.write_text(json.dumps(payload, ensure_ascii=True, indent=2) + "\n", encoding="utf-8") else: path.write_text(clean_content, encoding="utf-8") @@ -1147,6 +1474,225 @@ def _write_analysis_output( console.print(f"[green]✓ Wrote analysis output[/green] {path}") +def _is_ssh_uri(value: str) -> bool: + parsed = urlparse(value) + return parsed.scheme.lower() == "ssh" + + +def _is_http_url(value: str) -> bool: + parsed = urlparse(value) + return parsed.scheme.lower() in {"http", "https"} + + +def _download_text_url(url: str) -> str: + with urlopen(url, timeout=20) as resp: # noqa: S310 + charset = str(resp.headers.get_content_charset() or "utf-8") + raw: bytes = resp.read() + return raw.decode(charset, errors="replace") + + +def _looks_like_html(text: str) -> bool: + head = text.lstrip()[:256].lower() + return head.startswith(" bool: + stripped = text.strip() + if not stripped: + return False + if _looks_like_html(stripped): + return False + first_lines = stripped.splitlines()[:6] + markdown_markers = ("#", "- ", "* ", "```", ">", "---", "1. ") + return any(line.startswith(markdown_markers) for line in first_lines) + + +def _download_markdown_url(url: str) -> str: + """Download and validate a markdown payload from URL. + + Downloaded remote runbooks are never executed. They are persisted as plain + text and later parsed only for AI context retrieval. + """ + text = _download_text_url(url) + if not _looks_like_markdown(text): + raise ValueError(f"URL does not appear to be a Markdown payload: {url}") + return text + + +def _extract_markdown_links(index_html: str, base_url: str) -> list[str]: + pattern = r"href=[\"']([^\"']+\.md(?:\?[^\"']*)?)[\"']" + links = re.findall(pattern, index_html, flags=re.IGNORECASE) + resolved = [urljoin(base_url, link) for link in links] + seen: set[str] = set() + unique: list[str] = [] + for item in resolved: + if item not in seen: + seen.add(item) + unique.append(item) + return unique + + +def _build_ssh_source_config( + ssh_uri: str, + *, + identity_file: str | None, + jump_host: str | None, + ignore_ssh_config: bool, +) -> tuple[SSHConnectionConfig, str]: + """Build SSH connection config and remote path from an ssh:// URI.""" + parsed = urlparse(ssh_uri) + if parsed.scheme.lower() != "ssh": + raise ValueError(f"Not an ssh URI: {ssh_uri}") + if not parsed.hostname: + raise ValueError(f"SSH URI must include a host: {ssh_uri}") + if not parsed.path: + raise ValueError(f"SSH URI must include a remote path: {ssh_uri}") + + host = f"{parsed.username}@{parsed.hostname}" if parsed.username else parsed.hostname + remote_path = unquote(parsed.path) + + cfg = SSHConnectionConfig( + host=host, + port=parsed.port or 22, + identity_file=Path(identity_file).expanduser() if identity_file else None, + jump_host=jump_host, + ignore_ssh_config=ignore_ssh_config, + ) + return cfg, remote_path + + +def _run_ssh_source_command(config: SSHConnectionConfig, remote_command: str) -> str: + """Execute an SSH command for runbook source reads and return stdout.""" + client = SSHClient(config) + argv = client.build_ssh_argv(remote_command) + proc = subprocess.run(argv, capture_output=True, text=True, check=False) + if proc.returncode != 0: + details = (proc.stderr or proc.stdout or "unknown SSH error").strip() + raise RuntimeError(f"SSH source command failed: {details}") + return proc.stdout + + +def _materialize_runbooks_sync_path( + source_path: str, + *, + identity_file: str | None, + jump_host: str | None, + ignore_ssh_config: bool, +) -> tuple[Path, str, str | None]: + """Resolve local or remote runbooks directory for sync. + + Returns (local_dir, source_label, temp_dir_to_cleanup). + """ + if _is_http_url(source_path): + temp_dir = tempfile.mkdtemp(prefix="tai-runbooks-http-") + local_dir = Path(temp_dir) + source_url = source_path + parsed = urlparse(source_url) + + markdown_urls: list[str] + if parsed.path.lower().endswith(".md"): + markdown_urls = [source_url] + else: + index_html = _download_text_url(source_url) + markdown_urls = _extract_markdown_links(index_html, source_url) + + accepted = 0 + for idx, markdown_url in enumerate(markdown_urls, start=1): + try: + content = _download_markdown_url(markdown_url) + except ValueError: + # Ignore non-markdown payloads (e.g. HTML wrappers) even when + # linked with a .md suffix. + continue + url_path = urlparse(markdown_url).path + name = Path(unquote(url_path)).name or f"runbook-{idx}.md" + if not name.lower().endswith(".md"): + name = f"{name}.md" + (local_dir / name).write_text(content, encoding="utf-8") + accepted += 1 + + if accepted == 0: + raise ValueError( + "No valid Markdown runbooks found at web source. " + "Only markdown payloads are accepted." + ) + + return local_dir, source_url, temp_dir + + if not _is_ssh_uri(source_path): + local_dir = Path(source_path).expanduser().resolve() + return local_dir, str(local_dir), None + + config, remote_dir = _build_ssh_source_config( + source_path, + identity_file=identity_file, + jump_host=jump_host, + ignore_ssh_config=ignore_ssh_config, + ) + q_remote_dir = shlex.quote(remote_dir) + + ls_out = _run_ssh_source_command(config, f"ls -d {q_remote_dir}") + if not ls_out.strip(): + raise FileNotFoundError(f"Remote directory not found: {source_path}") + + find_out = _run_ssh_source_command( + config, + f"find {q_remote_dir} -maxdepth 1 -type f -name '*.md'", + ) + remote_files = sorted(line.strip() for line in find_out.splitlines() if line.strip()) + + temp_dir = tempfile.mkdtemp(prefix="tai-runbooks-") + local_dir = Path(temp_dir) + for remote_file in remote_files: + content = _run_ssh_source_command(config, f"cat {shlex.quote(remote_file)}") + (local_dir / Path(remote_file).name).write_text(content, encoding="utf-8") + + return local_dir, source_path, temp_dir + + +def _materialize_runbook_add_path( + source_file: str, + *, + identity_file: str | None, + jump_host: str | None, + ignore_ssh_config: bool, +) -> tuple[Path, str, str | None]: + """Resolve local or remote runbook file for add. + + Returns (local_file, source_label, temp_dir_to_cleanup). + """ + if _is_http_url(source_file): + parsed = urlparse(source_file) + if not parsed.path.lower().endswith(".md"): + raise ValueError("HTTP/HTTPS runbook source must point to a .md file") + content = _download_markdown_url(source_file) + temp_dir = tempfile.mkdtemp(prefix="tai-runbook-http-") + url_path = urlparse(source_file).path + name = Path(unquote(url_path)).name or "runbook.md" + if not name.lower().endswith(".md"): + name = f"{name}.md" + local_file = Path(temp_dir) / name + local_file.write_text(content, encoding="utf-8") + return local_file, source_file, temp_dir + + if not _is_ssh_uri(source_file): + local_file = Path(source_file).expanduser().resolve() + return local_file, str(local_file), None + + config, remote_file = _build_ssh_source_config( + source_file, + identity_file=identity_file, + jump_host=jump_host, + ignore_ssh_config=ignore_ssh_config, + ) + content = _run_ssh_source_command(config, f"cat {shlex.quote(remote_file)}") + + temp_dir = tempfile.mkdtemp(prefix="tai-runbook-") + local_file = Path(temp_dir) / Path(remote_file).name + local_file.write_text(content, encoding="utf-8") + return local_file, source_file, temp_dir + + # --------------------------------------------------------------------------- # runbooks sub-app # --------------------------------------------------------------------------- @@ -1156,12 +1702,30 @@ def _write_analysis_output( def runbooks_sync( path: Annotated[ str, - typer.Option("--path", help="Directory containing runbook Markdown files."), + typer.Option( + "--path", + help="Runbook source: local dir, ssh://host/path dir, or http(s) webroot URL.", + ), ] = "./runbooks", store_path: Annotated[ str, - typer.Option("--store", help="ChromaDB store path. Defaults to ~/.tai/runbooks."), + typer.Option( + "--store", + help="ChromaDB store path or remote URL. Defaults to ~/.tai/runbooks.", + ), ] = "~/.tai/runbooks", + store_user: Annotated[ + str | None, + typer.Option("--store-user", help="Runbook store login/user for remote ChromaDB URLs."), + ] = None, + store_password: Annotated[ + str | None, + typer.Option("--store-password", help="Runbook store password for remote ChromaDB URLs."), + ] = None, + env_file: Annotated[ + str | None, + typer.Option("--env-file", help="Optional dotenv file to resolve DB credentials."), + ] = None, ai_host: Annotated[ str, typer.Option("--ai-host", help="OpenAI-compatible AI backend URL."), @@ -1174,11 +1738,34 @@ def runbooks_sync( str, typer.Option("--ai-key", help="API key for the AI backend."), ] = "ollama", + identity_file: Annotated[ + str | None, + typer.Option("--identity-file", help="SSH private key path (for ssh:// --path)."), + ] = None, + jump_host: Annotated[ + str | None, + typer.Option("--jump-host", help="SSH bastion/jump host (for ssh:// --path)."), + ] = None, + ignore_ssh_config: Annotated[ + bool, + typer.Option("--ignore-ssh-config", help="Ignore ~/.ssh/config for ssh:// --path."), + ] = False, ) -> None: """Embed and index all runbooks from PATH into the persistent store.""" - from pathlib import Path + import shutil + + temp_dir: str | None = None + try: + runbooks_dir, source_label, temp_dir = _materialize_runbooks_sync_path( + path, + identity_file=identity_file, + jump_host=jump_host, + ignore_ssh_config=ignore_ssh_config, + ) + except Exception as exc: # noqa: BLE001 + console.print(f"[red]Runbooks source unavailable:[/red] {exc}") + raise typer.Exit(code=1) from exc - runbooks_dir = Path(path).expanduser().resolve() if not runbooks_dir.is_dir(): console.print(f"[red]Directory not found:[/red] {runbooks_dir}") raise typer.Exit(code=1) @@ -1186,25 +1773,75 @@ def runbooks_sync( ai_config = AIConfig(host=ai_host, model="", api_key=ai_key, embed_model=embed_model) ai = AIClient(ai_config) + env_values = _load_env_file(env_file) + resolved_store_user = _resolve_secret( + store_user, "TAI_RUNBOOK_STORE_USER", env_values + ) + resolved_store_password = _resolve_secret( + store_password, "TAI_RUNBOOK_STORE_PASSWORD", env_values + ) + display_store_path = _redact_url_password(store_path) + try: - store = RunbookStore(store_path) + store = RunbookStore( + store_path, + username=resolved_store_user, + password=resolved_store_password, + ) count = store.sync(runbooks_dir, ai) - console.print(f"[green]✓ Synced {count} runbook(s)[/green] → {store_path}") + msg = ( + f"[green]✓ Synced {count} runbook(s)[/green] from {source_label} " + f"→ {display_store_path}" + ) + console.print(msg) except Exception as exc: # noqa: BLE001 console.print(f"[red]Sync failed:[/red] {exc}") raise typer.Exit(code=1) from exc + finally: + if temp_dir: + shutil.rmtree(temp_dir, ignore_errors=True) @runbooks_app.command("list") def runbooks_list( store_path: Annotated[ str, - typer.Option("--store", help="ChromaDB store path. Defaults to ~/.tai/runbooks."), + typer.Option( + "--store", + help="ChromaDB store path or remote URL. Defaults to ~/.tai/runbooks.", + ), ] = "~/.tai/runbooks", + store_user: Annotated[ + str | None, + typer.Option("--store-user", help="Runbook store login/user for remote ChromaDB URLs."), + ] = None, + store_password: Annotated[ + str | None, + typer.Option( + "--store-password", + help="Runbook store password for remote ChromaDB URLs.", + ), + ] = None, + env_file: Annotated[ + str | None, + typer.Option("--env-file", help="Optional dotenv file to resolve DB credentials."), + ] = None, ) -> None: """List all indexed runbooks and their metadata.""" + env_values = _load_env_file(env_file) + resolved_store_user = _resolve_secret( + store_user, "TAI_RUNBOOK_STORE_USER", env_values + ) + resolved_store_password = _resolve_secret( + store_password, "TAI_RUNBOOK_STORE_PASSWORD", env_values + ) + try: - store = RunbookStore(store_path) + store = RunbookStore( + store_path, + username=resolved_store_user, + password=resolved_store_password, + ) entries = store.list_indexed() except Exception as exc: # noqa: BLE001 console.print(f"[red]Could not open store:[/red] {exc}") @@ -1224,11 +1861,32 @@ def runbooks_list( @runbooks_app.command("add") def runbooks_add( - file: Annotated[str, typer.Argument(help="Path to a single runbook Markdown file.")], + file: Annotated[ + str, + typer.Argument( + help="Runbook file source: local path, ssh://host/path/file.md, " + "or http(s) URL." + ), + ], store_path: Annotated[ str, - typer.Option("--store", help="ChromaDB store path. Defaults to ~/.tai/runbooks."), + typer.Option( + "--store", + help="ChromaDB store path or remote URL. Defaults to ~/.tai/runbooks.", + ), ] = "~/.tai/runbooks", + store_user: Annotated[ + str | None, + typer.Option("--store-user", help="Runbook store login/user for remote ChromaDB URLs."), + ] = None, + store_password: Annotated[ + str | None, + typer.Option("--store-password", help="Runbook store password for remote ChromaDB URLs."), + ] = None, + env_file: Annotated[ + str | None, + typer.Option("--env-file", help="Optional dotenv file to resolve DB credentials."), + ] = None, ai_host: Annotated[ str, typer.Option("--ai-host", help="OpenAI-compatible AI backend URL."), @@ -1241,11 +1899,34 @@ def runbooks_add( str, typer.Option("--ai-key", help="API key for the AI backend."), ] = "ollama", + identity_file: Annotated[ + str | None, + typer.Option("--identity-file", help="SSH private key path (for ssh:// file)."), + ] = None, + jump_host: Annotated[ + str | None, + typer.Option("--jump-host", help="SSH bastion/jump host (for ssh:// file)."), + ] = None, + ignore_ssh_config: Annotated[ + bool, + typer.Option("--ignore-ssh-config", help="Ignore ~/.ssh/config for ssh:// file."), + ] = False, ) -> None: """Embed and index a single runbook file into the persistent store.""" - from pathlib import Path + import shutil + + temp_dir: str | None = None + try: + runbook_path, source_label, temp_dir = _materialize_runbook_add_path( + file, + identity_file=identity_file, + jump_host=jump_host, + ignore_ssh_config=ignore_ssh_config, + ) + except Exception as exc: # noqa: BLE001 + console.print(f"[red]Runbook source unavailable:[/red] {exc}") + raise typer.Exit(code=1) from exc - runbook_path = Path(file).expanduser().resolve() if not runbook_path.is_file(): console.print(f"[red]File not found:[/red] {runbook_path}") raise typer.Exit(code=1) @@ -1253,13 +1934,33 @@ def runbooks_add( ai_config = AIConfig(host=ai_host, model="", api_key=ai_key, embed_model=embed_model) ai = AIClient(ai_config) + env_values = _load_env_file(env_file) + resolved_store_user = _resolve_secret( + store_user, "TAI_RUNBOOK_STORE_USER", env_values + ) + resolved_store_password = _resolve_secret( + store_password, "TAI_RUNBOOK_STORE_PASSWORD", env_values + ) + display_store_path = _redact_url_password(store_path) + try: - store = RunbookStore(store_path) + store = RunbookStore( + store_path, + username=resolved_store_user, + password=resolved_store_password, + ) store.sync_single(runbook_path, ai) - console.print(f"[green]✓ Indexed[/green] {runbook_path.name} → {store_path}") + msg = ( + f"[green]✓ Indexed[/green] {runbook_path.name} from {source_label} " + f"→ {display_store_path}" + ) + console.print(msg) except Exception as exc: # noqa: BLE001 console.print(f"[red]Add failed:[/red] {exc}") raise typer.Exit(code=1) from exc + finally: + if temp_dir: + shutil.rmtree(temp_dir, ignore_errors=True) def main() -> None: diff --git a/src/tai/history_store.py b/src/tai/history_store.py new file mode 100644 index 0000000..505bd63 --- /dev/null +++ b/src/tai/history_store.py @@ -0,0 +1,372 @@ +"""Persistent run history store backed by SQLite. + +Stores full per-run JSON payloads and allows retrieving host-specific history +to ground future analyses. +""" + +from __future__ import annotations + +import json +import sqlite3 +from contextlib import contextmanager +from dataclasses import dataclass +from pathlib import Path +from typing import Any +from urllib.parse import urlparse + +from tai.session_store import PastSession + +DEFAULT_HISTORY_DB_PATH = "~/.tai/history.db" + + +@dataclass(slots=True) +class HistoryRecord: + """Full history record persisted for one run.""" + + generated_at: str + issue: str + host: str + model: str + collection_total: int | None + collection_failed: int | None + collection_succeeded: int | None + prompt_tokens: int | None + completion_tokens: int | None + total_tokens: int | None + analysis: str + + +class RunHistoryStore: + """History store for host-scoped run payloads. + + Supported backends: + - SQLite local path (default, e.g. ``~/.tai/history.db``) + - SQLite URL (e.g. ``sqlite:////tmp/history.db``) + - PostgreSQL DSN (e.g. ``postgresql://user:pass@host:5432/dbname``) + """ + + def __init__(self, db_path: str | Path = DEFAULT_HISTORY_DB_PATH) -> None: + self._backend = "sqlite" + self._postgres_dsn: str | None = None + self._path: Path | None = None + + raw = str(db_path) + parsed = urlparse(raw) + if parsed.scheme in {"postgres", "postgresql"}: + self._backend = "postgres" + self._postgres_dsn = raw + elif parsed.scheme == "sqlite": + sqlite_path = parsed.path or "" + if sqlite_path.startswith("//"): + sqlite_path = sqlite_path[1:] + self._path = Path(sqlite_path).expanduser().resolve() + self._path.parent.mkdir(parents=True, exist_ok=True) + else: + self._path = Path(raw).expanduser().resolve() + self._path.parent.mkdir(parents=True, exist_ok=True) + + self._init_schema() + + def _connect(self) -> sqlite3.Connection: + if self._path is None: + raise RuntimeError("SQLite path is not configured for this history backend") + conn = sqlite3.connect(str(self._path)) + conn.row_factory = sqlite3.Row + return conn + + @contextmanager + def _connect_postgres(self) -> Any: + if self._postgres_dsn is None: + raise RuntimeError("PostgreSQL DSN is not configured for this history backend") + try: + import psycopg # type: ignore[import-not-found] + from psycopg.rows import dict_row # type: ignore[import-not-found] + except Exception as exc: # noqa: BLE001 + raise RuntimeError( + "PostgreSQL history backend requires psycopg. " + "Install with: pip install psycopg[binary]" + ) from exc + + conn = psycopg.connect(self._postgres_dsn, row_factory=dict_row) + try: + yield conn + finally: + conn.close() + + def _init_schema(self) -> None: + if self._backend == "postgres": + with self._connect_postgres() as conn: + with conn.cursor() as cur: + cur.execute( + """ + CREATE TABLE IF NOT EXISTS run_history ( + id BIGSERIAL PRIMARY KEY, + generated_at TEXT NOT NULL, + host TEXT NOT NULL, + issue TEXT NOT NULL, + model TEXT NOT NULL, + analysis TEXT NOT NULL, + payload_json TEXT NOT NULL + ) + """ + ) + cur.execute( + """ + CREATE INDEX IF NOT EXISTS idx_run_history_host_ts + ON run_history(host, generated_at DESC) + """ + ) + conn.commit() + return + + with self._connect() as conn: + conn.execute( + """ + CREATE TABLE IF NOT EXISTS run_history ( + id INTEGER PRIMARY KEY AUTOINCREMENT, + generated_at TEXT NOT NULL, + host TEXT NOT NULL, + issue TEXT NOT NULL, + model TEXT NOT NULL, + analysis TEXT NOT NULL, + payload_json TEXT NOT NULL + ) + """ + ) + conn.execute( + """ + CREATE INDEX IF NOT EXISTS idx_run_history_host_ts + ON run_history(host, generated_at DESC) + """ + ) + + def count(self, *, host: str | None = None) -> int: + if self._backend == "postgres": + with self._connect_postgres() as conn: + with conn.cursor() as cur: + if host is None: + cur.execute("SELECT COUNT(*) AS c FROM run_history") + else: + cur.execute( + "SELECT COUNT(*) AS c FROM run_history WHERE lower(host)=lower(%s)", + (host,), + ) + row = cur.fetchone() + if not row: + return 0 + return int(row["c"]) + + with self._connect() as conn: + if host is None: + row = conn.execute("SELECT COUNT(*) AS c FROM run_history").fetchone() + else: + row = conn.execute( + "SELECT COUNT(*) AS c FROM run_history WHERE lower(host)=lower(?)", + (host,), + ).fetchone() + count_value = row["c"] if row else 0 + return int(count_value) if isinstance(count_value, (int, float)) else 0 + + def add_payload(self, payload: dict[str, object]) -> int: + generated_at = str(payload.get("generated_at", "")) + host = str(payload.get("host", "")) + issue = str(payload.get("issue", "")) + model = str(payload.get("model", "")) + analysis = str(payload.get("analysis", "")) + payload_json = json.dumps(payload, ensure_ascii=True) + + if self._backend == "postgres": + with self._connect_postgres() as conn: + with conn.cursor() as cur: + cur.execute( + """ + INSERT INTO run_history( + generated_at, host, issue, model, analysis, + payload_json + ) + VALUES (%s, %s, %s, %s, %s, %s) + RETURNING id + """, + (generated_at, host, issue, model, analysis, payload_json), + ) + row = cur.fetchone() + conn.commit() + row_id = row["id"] if row else None + return int(row_id) if row_id is not None else 0 + + with self._connect() as conn: + cursor = conn.execute( + """ + INSERT INTO run_history(generated_at, host, issue, model, analysis, payload_json) + VALUES (?, ?, ?, ?, ?, ?) + """, + (generated_at, host, issue, model, analysis, payload_json), + ) + return cursor.lastrowid if cursor.lastrowid is not None else 0 + + def list_host_sessions(self, host: str, *, limit: int = 5) -> list[PastSession]: + return self.list_recent(host=host, limit=limit) + + def list_recent(self, *, host: str | None = None, limit: int = 20) -> list[PastSession]: + """Return recent records, optionally filtered by host.""" + if limit < 1: + return [] + + if self._backend == "postgres": + with self._connect_postgres() as conn: + with conn.cursor() as cur: + if host: + cur.execute( + """ + SELECT id, host, issue, analysis + FROM run_history + WHERE lower(host)=lower(%s) + ORDER BY generated_at DESC + LIMIT %s + """, + (host, limit), + ) + else: + cur.execute( + """ + SELECT id, host, issue, analysis + FROM run_history + ORDER BY generated_at DESC + LIMIT %s + """, + (limit,), + ) + rows = cur.fetchall() + + return [ + PastSession( + session_id=f"db-{row['id']}", + host=str(row["host"]), + issue=str(row["issue"]), + summary=str(row["analysis"]), + ) + for row in rows + ] + + with self._connect() as conn: + if host: + rows = conn.execute( + """ + SELECT id, host, issue, analysis + FROM run_history + WHERE lower(host)=lower(?) + ORDER BY generated_at DESC + LIMIT ? + """, + (host, limit), + ).fetchall() + else: + rows = conn.execute( + """ + SELECT id, host, issue, analysis + FROM run_history + ORDER BY generated_at DESC + LIMIT ? + """, + (limit,), + ).fetchall() + + return [ + PastSession( + session_id=f"db-{row['id']}", + host=str(row["host"]), + issue=str(row["issue"]), + summary=str(row["analysis"]), + ) + for row in rows + ] + + def search_keyword( + self, + keyword: str, + *, + host: str | None = None, + limit: int = 20, + ) -> list[PastSession]: + """Search issue/analysis text for *keyword*, optionally scoped by host.""" + term = keyword.strip() + if not term: + return self.list_recent(host=host, limit=limit) + + if limit < 1: + return [] + + like_pattern = f"%{term}%" + + if self._backend == "postgres": + with self._connect_postgres() as conn: + with conn.cursor() as cur: + if host: + cur.execute( + """ + SELECT id, host, issue, analysis + FROM run_history + WHERE lower(host)=lower(%s) + AND (issue ILIKE %s OR analysis ILIKE %s) + ORDER BY generated_at DESC + LIMIT %s + """, + (host, like_pattern, like_pattern, limit), + ) + else: + cur.execute( + """ + SELECT id, host, issue, analysis + FROM run_history + WHERE issue ILIKE %s OR analysis ILIKE %s + ORDER BY generated_at DESC + LIMIT %s + """, + (like_pattern, like_pattern, limit), + ) + rows = cur.fetchall() + + return [ + PastSession( + session_id=f"db-{row['id']}", + host=str(row["host"]), + issue=str(row["issue"]), + summary=str(row["analysis"]), + ) + for row in rows + ] + + with self._connect() as conn: + if host: + rows = conn.execute( + """ + SELECT id, host, issue, analysis + FROM run_history + WHERE lower(host)=lower(?) + AND (issue LIKE ? COLLATE NOCASE OR analysis LIKE ? COLLATE NOCASE) + ORDER BY generated_at DESC + LIMIT ? + """, + (host, like_pattern, like_pattern, limit), + ).fetchall() + else: + rows = conn.execute( + """ + SELECT id, host, issue, analysis + FROM run_history + WHERE issue LIKE ? COLLATE NOCASE OR analysis LIKE ? COLLATE NOCASE + ORDER BY generated_at DESC + LIMIT ? + """, + (like_pattern, like_pattern, limit), + ).fetchall() + + return [ + PastSession( + session_id=f"db-{row['id']}", + host=str(row["host"]), + issue=str(row["issue"]), + summary=str(row["analysis"]), + ) + for row in rows + ] diff --git a/src/tai/runbook_store.py b/src/tai/runbook_store.py index 4f34969..83a470f 100644 --- a/src/tai/runbook_store.py +++ b/src/tai/runbook_store.py @@ -14,10 +14,12 @@ Typical flow from __future__ import annotations +import base64 import re from dataclasses import dataclass, field from pathlib import Path from typing import TYPE_CHECKING, Any +from urllib.parse import urlparse if TYPE_CHECKING: from tai.ai_client import AIClient @@ -100,11 +102,19 @@ class RunbookStore: Defaults to ``~/.tai/runbooks``. """ - def __init__(self, store_path: str | Path = DEFAULT_STORE_PATH) -> None: + def __init__( + self, + store_path: str | Path = DEFAULT_STORE_PATH, + *, + username: str | None = None, + password: str | None = None, + ) -> None: import chromadb # optional dep — imported lazily - path = Path(store_path).expanduser().resolve() - path.mkdir(parents=True, exist_ok=True) + raw_store = str(store_path) + parsed = urlparse(raw_store) + is_remote = parsed.scheme in {"http", "https"} + settings = None try: from chromadb.config import Settings @@ -119,10 +129,28 @@ class RunbookStore: # does not expose the real config module. settings = None - if settings is None: - self._client = chromadb.PersistentClient(path=str(path)) + if is_remote: + host = parsed.hostname or "localhost" + port = parsed.port or (443 if parsed.scheme == "https" else 80) + ssl = parsed.scheme == "https" + auth_user = username if username is not None else parsed.username + auth_pass = password if password is not None else parsed.password + headers = None + if auth_user is not None and auth_pass is not None: + token = base64.b64encode(f"{auth_user}:{auth_pass}".encode()).decode("ascii") + headers = {"Authorization": f"Basic {token}"} + + if headers is None: + self._client = chromadb.HttpClient(host=host, port=port, ssl=ssl) + else: + self._client = chromadb.HttpClient(host=host, port=port, ssl=ssl, headers=headers) else: - self._client = chromadb.PersistentClient(path=str(path), settings=settings) + path = Path(store_path).expanduser().resolve() + path.mkdir(parents=True, exist_ok=True) + if settings is None: + self._client = chromadb.PersistentClient(path=str(path)) + else: + self._client = chromadb.PersistentClient(path=str(path), settings=settings) self._collection: Any = self._client.get_or_create_collection( name=_COLLECTION_NAME, metadata={"hnsw:space": "cosine"}, diff --git a/tai-live-ai-check.md b/tai-live-ai-check.md new file mode 100644 index 0000000..0bff022 --- /dev/null +++ b/tai-live-ai-check.md @@ -0,0 +1,3 @@ +# Live AI check + +This verifies real embedding calls. diff --git a/tests/test_cli.py b/tests/test_cli.py index c68b3af..fc7fafe 100644 --- a/tests/test_cli.py +++ b/tests/test_cli.py @@ -1,11 +1,22 @@ import json +import os +import re from pathlib import Path from types import SimpleNamespace from unittest.mock import AsyncMock, MagicMock +import pytest from typer.testing import CliRunner -from tai.cli import app +from tai.cli import ( + _download_markdown_url, + _inject_url_credentials, + _load_env_file, + _materialize_runbook_add_path, + _materialize_runbooks_sync_path, + _resolve_secret, + app, +) from tai.collectors import CollectedItem, CollectionReport from tai.rag_retriever import Chunk, EmbeddedChunk from tai.ssh_client import SSHCommandResult @@ -344,7 +355,7 @@ def test_interactive_rag_debug_prints_retrieval_scores(monkeypatch) -> None: # def test_history_command_lists_sessions(monkeypatch) -> None: # type: ignore[no-untyped-def] class FakeStore: - def __init__(self, _path: str) -> None: + def __init__(self, _path: str, **_kwargs) -> None: pass def list_recent(self, *, host: str | None = None, limit: int = 20): @@ -360,12 +371,12 @@ def test_history_command_lists_sessions(monkeypatch) -> None: # type: ignore[no ] return [] - monkeypatch.setattr("tai.cli.SessionStore", FakeStore) + monkeypatch.setattr("tai.cli.RunHistoryStore", FakeStore) runner = CliRunner() result = runner.invoke( app, - ["history", "--session-memory", "~/.tai/sessions", "--host", "web01"], + ["history", "--history-db", "~/.tai/history.db", "--host", "web01"], ) assert result.exit_code == 0 @@ -375,7 +386,7 @@ def test_history_command_lists_sessions(monkeypatch) -> None: # type: ignore[no def test_history_command_exports_markdown(monkeypatch, tmp_path: Path) -> None: # type: ignore[no-untyped-def] class FakeStore: - def __init__(self, _path: str) -> None: + def __init__(self, _path: str, **_kwargs) -> None: pass def list_recent(self, *, host: str | None = None, limit: int = 20): @@ -389,13 +400,13 @@ def test_history_command_exports_markdown(monkeypatch, tmp_path: Path) -> None: ) ] - monkeypatch.setattr("tai.cli.SessionStore", FakeStore) + monkeypatch.setattr("tai.cli.RunHistoryStore", FakeStore) export_path = tmp_path / "history.md" runner = CliRunner() result = runner.invoke( app, - ["history", "--session-memory", "~/.tai/sessions", "--export", str(export_path)], + ["history", "--history-db", "~/.tai/history.db", "--export", str(export_path)], ) assert result.exit_code == 0 @@ -440,11 +451,12 @@ def test_interactive_history_without_store_shows_hint(monkeypatch) -> None: # t "5566", "--no-probe", "--interactive", + "--no-history", ], ) assert result.exit_code == 0 - assert "Session memory is disabled" in result.stdout + assert "History DB is disabled" in result.stdout def test_run_analyze_writes_output_file(monkeypatch, tmp_path: Path) -> None: # type: ignore[no-untyped-def] @@ -467,9 +479,10 @@ def test_run_analyze_writes_output_file(monkeypatch, tmp_path: Path) -> None: # ) monkeypatch.setattr("tai.cli.collect_from_plan", fake_collect_from_plan) + response = SimpleNamespace(content="Root Cause\n\nEvidence\n\nRecommended Actions") monkeypatch.setattr( "tai.cli.AIClient.complete", - lambda *_args, **_kwargs: SimpleNamespace(content="Root Cause\n\nEvidence\n\nRecommended Actions"), + lambda *_args, **_kwargs: response, ) output_path = tmp_path / "analysis.md" @@ -543,7 +556,304 @@ def test_run_analyze_writes_json_output_and_strips_ansi(monkeypatch, tmp_path: P assert result.exit_code == 0 payload = json.loads(output_path.read_text(encoding="utf-8")) + assert payload["schema"] == "tai.analysis.v1" + assert "generated_at" in payload assert payload["issue"] == "apache failed" assert payload["host"] == "ssh.archflux.net" + assert payload["collection"] == {"total": 1, "failed": 0, "succeeded": 1} + assert payload["token_usage"] == { + "prompt_tokens": None, + "completion_tokens": None, + "total_tokens": None, + } assert "Root Cause" in payload["analysis"] assert "\u001b" not in payload["analysis"] + + +def test_run_analyze_writes_history_db_record(monkeypatch, tmp_path: Path) -> None: # type: ignore[no-untyped-def] + _mock_session(monkeypatch) + + async def fake_collect_from_plan(_session, _plan) -> CollectionReport: # type: ignore[no-untyped-def] + return CollectionReport( + host="ssh.archflux.net", + items=[ + CollectedItem( + name="kernel", + result=SSHCommandResult( + command="uname -a", + exit_code=0, + stdout="Linux test", + stderr="", + ), + ), + ], + ) + + monkeypatch.setattr("tai.cli.collect_from_plan", fake_collect_from_plan) + response = SimpleNamespace(content="Root Cause\n\nEvidence\n\nRecommended Actions") + monkeypatch.setattr( + "tai.cli.AIClient.complete", + lambda *_args, **_kwargs: response, + ) + + history_db = tmp_path / "history.db" + runner = CliRunner() + result = runner.invoke( + app, + [ + "run", "apache failed", + "--host", + "ssh.archflux.net", + "--port", + "5566", + "--no-probe", + "--analyze", + "--history-db", + str(history_db), + ], + ) + + assert result.exit_code == 0 + + import sqlite3 + + with sqlite3.connect(str(history_db)) as conn: + row = conn.execute( + "SELECT host, issue, payload_json FROM run_history ORDER BY id DESC LIMIT 1" + ).fetchone() + + assert row is not None + assert row[0] == "ssh.archflux.net" + assert row[1] == "apache failed" + payload = json.loads(row[2]) + assert payload["schema"] == "tai.analysis.v1" + assert payload["host"] == "ssh.archflux.net" + + +def test_materialize_runbooks_sync_path_http_webroot(monkeypatch, tmp_path: Path) -> None: # type: ignore[no-untyped-def] + html = 'nginxssh' + + def fake_download(url: str) -> str: + if url == "https://kb.example/runbooks/": + return html + if url.endswith("nginx.md"): + return "---\nservice: nginx\n---\nbody" + if url.endswith("ssh.md"): + return "---\nservice: ssh\n---\nbody" + raise AssertionError(url) + + monkeypatch.setattr("tai.cli._download_text_url", fake_download) + source_dir, label, temp_dir = _materialize_runbooks_sync_path( + "https://kb.example/runbooks/", + identity_file=None, + jump_host=None, + ignore_ssh_config=False, + ) + + assert label == "https://kb.example/runbooks/" + assert temp_dir is not None + assert (source_dir / "nginx.md").is_file() + assert (source_dir / "ssh.md").is_file() + + +def test_materialize_runbook_add_path_http_url(monkeypatch) -> None: # type: ignore[no-untyped-def] + monkeypatch.setattr( + "tai.cli._download_markdown_url", + lambda _url: "---\nservice: nginx\n---\nbody", + ) + source_file, label, temp_dir = _materialize_runbook_add_path( + "https://kb.example/runbooks/nginx.md", + identity_file=None, + jump_host=None, + ignore_ssh_config=False, + ) + + assert label == "https://kb.example/runbooks/nginx.md" + assert temp_dir is not None + assert source_file.name == "nginx.md" + assert source_file.read_text(encoding="utf-8").startswith("---") + + +def test_download_markdown_url_rejects_html(monkeypatch) -> None: # type: ignore[no-untyped-def] + monkeypatch.setattr( + "tai.cli._download_text_url", + lambda _url: "not markdown", + ) + + with pytest.raises(ValueError, match="does not appear to be a Markdown payload"): + _download_markdown_url("https://kb.example/runbooks/nginx.md") + + +def test_materialize_runbooks_sync_path_http_skips_html_wrappers(monkeypatch) -> None: # type: ignore[no-untyped-def] + html = 'nginxssh' + + def fake_download(url: str) -> str: + if url == "https://kb.example/runbooks/": + return html + if url.endswith("nginx.md"): + return "---\nservice: nginx\n---\nbody" + if url.endswith("ssh.md"): + return "wrapper" + raise AssertionError(url) + + monkeypatch.setattr("tai.cli._download_text_url", fake_download) + source_dir, _label, temp_dir = _materialize_runbooks_sync_path( + "https://kb.example/runbooks/", + identity_file=None, + jump_host=None, + ignore_ssh_config=False, + ) + + assert temp_dir is not None + assert (source_dir / "nginx.md").is_file() + assert not (source_dir / "ssh.md").exists() + + +def test_materialize_runbook_add_path_http_requires_md_suffix() -> None: + with pytest.raises(ValueError, match="must point to a .md file"): + _materialize_runbook_add_path( + "https://kb.example/runbooks/", + identity_file=None, + jump_host=None, + ignore_ssh_config=False, + ) + + +def test_runbooks_sync_accepts_ssh_source(monkeypatch, tmp_path: Path) -> None: # type: ignore[no-untyped-def] + runbooks_dir = tmp_path / "remote-runbooks" + runbooks_dir.mkdir(parents=True) + (runbooks_dir / "nginx.md").write_text("---\nservice: nginx\n---\nbody", encoding="utf-8") + + monkeypatch.setattr( + "tai.cli._materialize_runbooks_sync_path", + lambda *_args, **_kwargs: (runbooks_dir, "ssh://ops@host/runbooks", None), + ) + + class FakeStore: + def __init__(self, _path: str, **_kwargs) -> None: + pass + + def sync(self, _dir: Path, _ai): + return 1 + + monkeypatch.setattr("tai.cli.RunbookStore", FakeStore) + monkeypatch.setattr("tai.cli.AIClient", lambda *_a, **_k: object()) + + runner = CliRunner() + result = runner.invoke( + app, + [ + "runbooks", + "sync", + "--path", + "ssh://ops@host/runbooks", + "--store", + "~/.tai/runbooks", + ], + ) + + assert result.exit_code == 0 + assert "Synced 1 runbook(s)" in result.stdout + assert "ssh://ops@host/runbooks" in result.stdout + + +def test_runbooks_add_accepts_https_source(monkeypatch) -> None: # type: ignore[no-untyped-def] + import tempfile + + fd, temp_name = tempfile.mkstemp(prefix="tai-runbook-test-", suffix=".md") + os.close(fd) + Path(temp_name).write_text("---\nservice: nginx\n---\nbody", encoding="utf-8") + + monkeypatch.setattr( + "tai.cli._materialize_runbook_add_path", + lambda *_args, **_kwargs: (Path(temp_name), "https://kb.example/nginx.md", None), + ) + + class FakeStore: + def __init__(self, _path: str, **_kwargs) -> None: + pass + + def sync_single(self, _path: Path, _ai): + return None + + monkeypatch.setattr("tai.cli.RunbookStore", FakeStore) + monkeypatch.setattr("tai.cli.AIClient", lambda *_a, **_k: object()) + + runner = CliRunner() + result = runner.invoke( + app, + [ + "runbooks", + "add", + "https://kb.example/nginx.md", + "--store", + "~/.tai/runbooks", + ], + ) + + assert result.exit_code == 0 + assert "Indexed" in result.stdout + assert "https://kb.example/nginx.md" in result.stdout + + Path(temp_name).unlink(missing_ok=True) + + +def test_inject_url_credentials_postgres() -> None: + target = "postgresql://db.example.com:5432/tai" + rendered = _inject_url_credentials( + target, + user="tai_user", + password="secret", + schemes={"postgresql", "postgres"}, + ) + assert rendered.startswith("postgresql://tai_user:secret@db.example.com:5432/tai") + + +def test_inject_url_credentials_ignores_non_matching_scheme() -> None: + target = "~/.tai/history.db" + rendered = _inject_url_credentials( + target, + user="tai_user", + password="secret", + schemes={"postgresql", "postgres"}, + ) + assert rendered == target + + +def test_load_env_file_and_resolve_secret(tmp_path: Path, monkeypatch) -> None: # type: ignore[no-untyped-def] + env_file = tmp_path / ".env" + env_file.write_text( + "TAI_HISTORY_DB_USER=from_file\n" + "TAI_HISTORY_DB_PASSWORD=from_file_pw\n", + encoding="utf-8", + ) + + values = _load_env_file(str(env_file)) + assert values["TAI_HISTORY_DB_USER"] == "from_file" + assert values["TAI_HISTORY_DB_PASSWORD"] == "from_file_pw" + + monkeypatch.setenv("TAI_HISTORY_DB_USER", "from_env") + assert _resolve_secret(None, "TAI_HISTORY_DB_USER", values) == "from_file" + assert _resolve_secret("from_cli", "TAI_HISTORY_DB_USER", values) == "from_cli" + + +def test_man_page_covers_cli_long_options() -> None: + runner = CliRunner() + help_invocations = [ + ["run", "--help"], + ["history", "--help"], + ["runbooks", "sync", "--help"], + ["runbooks", "list", "--help"], + ["runbooks", "add", "--help"], + ] + + documented = Path("docs/tai.1").read_text(encoding="utf-8") + discovered: set[str] = set() + for args in help_invocations: + result = runner.invoke(app, args) + assert result.exit_code == 0, f"help command failed for: {' '.join(args)}" + discovered.update(re.findall(r"--[a-z0-9][a-z0-9-]*", result.stdout)) + + discovered.discard("--help") + missing = sorted(option for option in discovered if option not in documented) + assert missing == [], f"Missing options in docs/tai.1: {', '.join(missing)}" diff --git a/tests/test_history_store.py b/tests/test_history_store.py new file mode 100644 index 0000000..c1297ad --- /dev/null +++ b/tests/test_history_store.py @@ -0,0 +1,115 @@ +"""Tests for SQLite-backed run history storage.""" + +from __future__ import annotations + +from pathlib import Path + +from tai.history_store import RunHistoryStore + + +def test_history_store_add_and_count(tmp_path) -> None: # type: ignore[no-untyped-def] + store = RunHistoryStore(tmp_path / "history.db") + assert store.count() == 0 + + payload = { + "schema": "tai.analysis.v1", + "generated_at": "2026-05-11T12:00:00+00:00", + "issue": "sshd failed", + "host": "ssh.archflux.net", + "model": "gemma3:4b", + "collection": {"total": 5, "failed": 1, "succeeded": 4}, + "token_usage": {"prompt_tokens": 10, "completion_tokens": 20, "total_tokens": 30}, + "analysis": "Root Cause...", + } + store.add_payload(payload) + + assert store.count() == 1 + assert store.count(host="ssh.archflux.net") == 1 + assert store.count(host="other") == 0 + + +def test_history_store_list_host_sessions(tmp_path) -> None: # type: ignore[no-untyped-def] + store = RunHistoryStore(tmp_path / "history.db") + store.add_payload( + { + "schema": "tai.analysis.v1", + "generated_at": "2026-05-11T12:00:00+00:00", + "issue": "issue one", + "host": "ssh.archflux.net", + "model": "gemma3:4b", + "collection": {"total": 1, "failed": 0, "succeeded": 1}, + "token_usage": {"prompt_tokens": 1, "completion_tokens": 2, "total_tokens": 3}, + "analysis": "first", + } + ) + store.add_payload( + { + "schema": "tai.analysis.v1", + "generated_at": "2026-05-11T12:05:00+00:00", + "issue": "issue two", + "host": "ssh.archflux.net", + "model": "gemma3:4b", + "collection": {"total": 1, "failed": 0, "succeeded": 1}, + "token_usage": {"prompt_tokens": 1, "completion_tokens": 2, "total_tokens": 3}, + "analysis": "second", + } + ) + + sessions = store.list_host_sessions("ssh.archflux.net", limit=2) + assert len(sessions) == 2 + assert sessions[0].issue == "issue two" + assert sessions[1].issue == "issue one" + + +def test_history_store_list_recent_and_search_keyword(tmp_path) -> None: # type: ignore[no-untyped-def] + store = RunHistoryStore(tmp_path / "history.db") + store.add_payload( + { + "schema": "tai.analysis.v1", + "generated_at": "2026-05-11T13:00:00+00:00", + "issue": "nginx failed", + "host": "web01", + "model": "gemma3:4b", + "collection": {"total": 1, "failed": 0, "succeeded": 1}, + "token_usage": {"prompt_tokens": 1, "completion_tokens": 2, "total_tokens": 3}, + "analysis": "nginx config typo", + } + ) + store.add_payload( + { + "schema": "tai.analysis.v1", + "generated_at": "2026-05-11T13:10:00+00:00", + "issue": "sshd failed", + "host": "ssh.archflux.net", + "model": "gemma3:4b", + "collection": {"total": 1, "failed": 0, "succeeded": 1}, + "token_usage": {"prompt_tokens": 1, "completion_tokens": 2, "total_tokens": 3}, + "analysis": "sshd key mismatch", + } + ) + + recent = store.list_recent(limit=2) + assert len(recent) == 2 + assert recent[0].issue == "sshd failed" + + matches = store.search_keyword("key", host="ssh.archflux.net", limit=5) + assert len(matches) == 1 + assert matches[0].host == "ssh.archflux.net" + + +def test_history_store_accepts_sqlite_url(tmp_path: Path) -> None: + db_file = tmp_path / "history-url.db" + store = RunHistoryStore(f"sqlite:///{db_file}") + store.add_payload( + { + "schema": "tai.analysis.v1", + "generated_at": "2026-05-11T13:20:00+00:00", + "issue": "test", + "host": "host1", + "model": "gemma3:4b", + "collection": {"total": 1, "failed": 0, "succeeded": 1}, + "token_usage": {"prompt_tokens": 1, "completion_tokens": 1, "total_tokens": 2}, + "analysis": "ok", + } + ) + assert store.count(host="host1") == 1 diff --git a/tests/test_runbook_store.py b/tests/test_runbook_store.py index a6afb7b..d752c45 100644 --- a/tests/test_runbook_store.py +++ b/tests/test_runbook_store.py @@ -100,6 +100,7 @@ def _make_chromadb_mock() -> MagicMock: client.get_or_create_collection.return_value = collection chroma_mod = MagicMock() chroma_mod.PersistentClient.return_value = client + chroma_mod.HttpClient.return_value = client return chroma_mod @@ -251,3 +252,28 @@ def test_runbook_store_sync_single_missing_file_raises(tmp_path: Path) -> None: store = RunbookStore(tmp_path / "store") with pytest.raises(FileNotFoundError): store.sync_single(tmp_path / "missing.md", ai) + + +def test_runbook_store_remote_url_uses_http_client() -> None: + chroma_mock = _make_chromadb_mock() + + with patch.dict("sys.modules", {"chromadb": chroma_mock}): + store = RunbookStore("https://chroma.example.com:8443") + assert store.count() == 0 + + chroma_mock.HttpClient.assert_called_once_with(host="chroma.example.com", port=8443, ssl=True) + + +def test_runbook_store_remote_url_uses_http_client_with_basic_auth() -> None: + chroma_mock = _make_chromadb_mock() + + with patch.dict("sys.modules", {"chromadb": chroma_mock}): + store = RunbookStore("https://chroma.example.com:8443", username="tai", password="secret") + assert store.count() == 0 + + args = chroma_mock.HttpClient.call_args.kwargs + assert args["host"] == "chroma.example.com" + assert args["port"] == 8443 + assert args["ssl"] is True + assert "headers" in args + assert str(args["headers"].get("Authorization", "")).startswith("Basic ")