update

tests/test_ai.py

@@ -5,6 +5,7 @@ from unittest.mock import MagicMock, patch
 from tai.ai_client import DEFAULT_AI_HOST, DEFAULT_MODEL, AIClient, AIConfig
 from tai.collectors import CollectedItem, CollectionReport
 from tai.prompt_builder import build_followup_message, build_system_prompt, build_user_message
+from tai.session_store import PastSession
 from tai.ssh_client import SSHCommandResult
 
 # ---------------------------------------------------------------------------
@@ -221,6 +222,24 @@ def test_build_user_message_handles_no_output() -> None:
     assert "no output" in msg
 
 
+def test_build_user_message_includes_prior_session_context() -> None:
+    report = _make_report([("kernel", "uname -a", 0, "Linux web01", "")])
+    msg = build_user_message(
+        "sssd broken",
+        report,
+        past_sessions=[
+            PastSession(
+                session_id="20260506T120000Z",
+                host="web01",
+                issue="sssd broken",
+                summary="Root cause was missing sssd package.",
+            )
+        ],
+    )
+    assert "Similar prior sessions" in msg
+    assert "missing sssd package" in msg
+
+
 def test_build_followup_message_includes_question_context() -> None:
     report = _make_report([("kernel", "uname -a", 0, "Linux web01", "")])
     msg = build_followup_message(

tests/test_plan.py

@@ -107,8 +107,12 @@ def test_sssd_in_issue_adds_presence_service_and_config_commands() -> None:
     assert "binary-sssd-1" in names
     assert "service-sssd" in names
     assert "journal-sssd" in names
+    assert "package-rpm-sssd-1" in names
+    assert "package-dpkg-sssd-1" in names
     assert any("cat /etc/sssd/sssd.conf" in c for c in cmds)
     assert any("ls -l /usr/sbin/sssd" in c for c in cmds)
+    assert any("rpm -q sssd" in c for c in cmds)
+    assert any("dpkg-query -W sssd" in c for c in cmds)
     assert any("list-unit-files sssd.service" in c for c in cmds)
 
 
@@ -119,8 +123,12 @@ def test_docker_presence_probe_checks_package_and_binary() -> None:
     assert "unit-file-docker" in names
     assert "binary-docker-1" in names
     assert "binary-docker-2" in names
+    assert "package-rpm-docker-1" in names
+    assert "package-dpkg-docker-1" in names
     assert any("ls -l /usr/bin/docker" in c for c in cmds)
     assert any("ls -l /usr/bin/dockerd" in c for c in cmds)
+    assert any("rpm -q docker" in c for c in cmds)
+    assert any("dpkg-query -W docker" in c for c in cmds)
 
 
 def test_unknown_service_name_no_config_cat() -> None:
@@ -183,6 +191,16 @@ def test_extract_services_case_insensitive() -> None:
     assert "nginx" in _extract_services("NGINX failed")
 
 
+def test_extract_services_detects_generic_service_name() -> None:
+    services = _extract_services("myweirdapp service keeps failing")
+    assert "myweirdapp" in services
+
+
+def test_extract_services_detects_dot_service_pattern() -> None:
+    services = _extract_services("please check foobar.service on this host")
+    assert "foobar" in services
+
+
 # ---------------------------------------------------------------------------
 # Plan length sanity
 # ---------------------------------------------------------------------------

tests/test_session_store.py

@@ -0,0 +1,79 @@
+"""Tests for session_store with mocked ChromaDB."""
+
+from __future__ import annotations
+
+from pathlib import Path
+from unittest.mock import MagicMock, patch
+
+from tai.session_store import PastSession, SessionStore, _build_embed_text
+
+
+def _make_chromadb_mock() -> MagicMock:
+    collection = MagicMock()
+    collection.count.return_value = 0
+    client = MagicMock()
+    client.get_or_create_collection.return_value = collection
+    chroma_mod = MagicMock()
+    chroma_mod.PersistentClient.return_value = client
+    return chroma_mod
+
+
+def _make_ai_mock(embedding: list[float] | None = None) -> MagicMock:
+    ai = MagicMock()
+    ai.embed.return_value = embedding or [0.1, 0.2, 0.3]
+    return ai
+
+
+def test_build_embed_text_contains_host_issue_and_summary() -> None:
+    text = _build_embed_text(host="web01", issue="sssd broken", summary="Unit missing")
+    assert "host: web01" in text
+    assert "issue: sssd broken" in text
+    assert "Unit missing" in text
+
+
+def test_index_session_upserts_with_metadata(tmp_path: Path) -> None:
+    chroma_mock = _make_chromadb_mock()
+    collection = chroma_mock.PersistentClient.return_value.get_or_create_collection.return_value
+    ai = _make_ai_mock()
+
+    with patch.dict("sys.modules", {"chromadb": chroma_mock}):
+        store = SessionStore(tmp_path / "store")
+        session_id = store.index_session("web01", "sssd broken", "summary text", ai)
+
+    assert session_id
+    collection.upsert.assert_called_once()
+    args = collection.upsert.call_args.kwargs
+    assert args["metadatas"][0]["host"] == "web01"
+    assert args["metadatas"][0]["issue"] == "sssd broken"
+
+
+def test_query_returns_empty_when_no_docs(tmp_path: Path) -> None:
+    chroma_mock = _make_chromadb_mock()
+    ai = _make_ai_mock()
+
+    with patch.dict("sys.modules", {"chromadb": chroma_mock}):
+        store = SessionStore(tmp_path / "store")
+        results = store.query("why sssd", "web01", ai)
+
+    assert results == []
+
+
+def test_query_returns_past_sessions(tmp_path: Path) -> None:
+    chroma_mock = _make_chromadb_mock()
+    collection = chroma_mock.PersistentClient.return_value.get_or_create_collection.return_value
+    collection.count.return_value = 1
+    collection.query.return_value = {
+        "ids": [["20260506T120000Z"]],
+        "documents": [["Root cause: package missing"]],
+        "metadatas": [[{"host": "web01", "issue": "sssd broken"}]],
+    }
+    ai = _make_ai_mock()
+
+    with patch.dict("sys.modules", {"chromadb": chroma_mock}):
+        store = SessionStore(tmp_path / "store")
+        results = store.query("sssd issue", "web01", ai)
+
+    assert len(results) == 1
+    assert isinstance(results[0], PastSession)
+    assert results[0].host == "web01"
+    assert "package missing" in results[0].summary

tests/test_ssh_client.py

@@ -82,6 +82,8 @@ def test_allows_expected_read_only_commands() -> None:
         "systemctl status apache2",
         "cat /etc/hosts",
         "ss -lntp",
+        "rpm -q sssd",
+        "dpkg-query -W sssd",
     ]:
         client.validate_read_only_command(command)