feat(workspaces): add unified /workspace lifecycle, target persistence, and workspace-scoped RAG

- Introduce  command for CLI and TUI with create/activate, list, info, note, clear, export, import, and help actions
- Persist workspace state via  marker and enriched  (targets, operator notes, last_active_at, last_target)
- Restore  on workspace activation and sync it to UI banner, agent state, and CLI output
- Enforce target normalization and ensure  always exists in workspace targets
- Route loot output to  when a workspace is active
- Prefer workspace-local knowledge paths for indexing and RAG resolution
- Persist RAG indexes per workspace and load existing indexes before re-indexing
- Add deterministic workspace export/import utilities (excluding caches)
- Integrate workspace handling into TUI slash commands with modal help screen
This commit is contained in:
giveen
2026-01-19 08:41:38 -07:00
parent 50c8ec1936
commit e8ab673a13
20 changed files with 1439 additions and 56 deletions

View File

@@ -0,0 +1,50 @@
import os
from pathlib import Path
import pytest
from pentestagent.workspaces.manager import WorkspaceManager
from pentestagent.knowledge.rag import RAGEngine
from pentestagent.knowledge.indexer import KnowledgeIndexer
def test_rag_and_indexer_use_workspace(tmp_path, monkeypatch):
# Use tmp_path as the project root
monkeypatch.chdir(tmp_path)
wm = WorkspaceManager(root=tmp_path)
name = "ws_test"
wm.create(name)
wm.set_active(name)
# Create a sample source file in the workspace sources
src_dir = tmp_path / "workspaces" / name / "knowledge" / "sources"
src_dir.mkdir(parents=True, exist_ok=True)
sample = src_dir / "sample.md"
sample.write_text("# Sample\n\nThis is a test knowledge document for RAG indexing.")
# Ensure KnowledgeIndexer picks up the workspace source when indexing default 'knowledge'
ki = KnowledgeIndexer()
docs, result = ki.index_directory(Path("knowledge"))
assert result.indexed_files >= 1
assert len(docs) >= 1
# Ensure the document source path points at the workspace file
assert any("workspaces" in d.source and "sample.md" in d.source for d in docs)
# Now run RAGEngine to build embeddings and verify saved index file appears
rag = RAGEngine(use_local_embeddings=True)
rag.index()
emb_path = tmp_path / "workspaces" / name / "knowledge" / "embeddings" / "index.pkl"
assert emb_path.exists(), f"Expected saved index at {emb_path}"
# Ensure RAG engine has documents/chunks loaded
assert rag.get_chunk_count() >= 1
assert rag.get_document_count() >= 1
# Now create a new RAGEngine and ensure it loads persisted index automatically
rag2 = RAGEngine(use_local_embeddings=True)
# If load-on-init doesn't run, calling index() should load from saved file
rag2.index()
assert rag2.get_chunk_count() >= 1

96
tests/test_workspace.py Normal file
View File

@@ -0,0 +1,96 @@
import os
from pathlib import Path
import pytest
from pentestagent.workspaces.manager import WorkspaceManager, WorkspaceError
def test_invalid_workspace_names(tmp_path: Path):
wm = WorkspaceManager(root=tmp_path)
bad_names = ["../escape", "name/with/slash", "..", ""]
# overlong name
bad_names.append("a" * 65)
for n in bad_names:
with pytest.raises(WorkspaceError):
wm.create(n)
def test_create_and_idempotent(tmp_path: Path):
wm = WorkspaceManager(root=tmp_path)
name = "eng1"
meta = wm.create(name)
assert (tmp_path / "workspaces" / name).exists()
assert (tmp_path / "workspaces" / name / "meta.yaml").exists()
# create again should not raise and should return meta
meta2 = wm.create(name)
assert meta2["name"] == name
def test_set_get_active(tmp_path: Path):
wm = WorkspaceManager(root=tmp_path)
name = "activews"
wm.create(name)
wm.set_active(name)
assert wm.get_active() == name
marker = tmp_path / "workspaces" / ".active"
assert marker.exists()
assert marker.read_text(encoding="utf-8").strip() == name
def test_add_list_remove_targets(tmp_path: Path):
wm = WorkspaceManager(root=tmp_path)
name = "targets"
wm.create(name)
added = wm.add_targets(name, ["192.168.1.1", "192.168.0.0/16", "Example.COM"]) # hostname mixed case
# normalized entries
assert "192.168.1.1" in added
assert "192.168.0.0/16" in added
assert "example.com" in added
# dedupe
added2 = wm.add_targets(name, ["192.168.1.1", "example.com"])
assert len(added2) == len(added)
# remove
after = wm.remove_target(name, "192.168.1.1")
assert "192.168.1.1" not in after
def test_persistence_across_instances(tmp_path: Path):
wm1 = WorkspaceManager(root=tmp_path)
name = "persist"
wm1.create(name)
wm1.add_targets(name, ["10.0.0.1", "host.local"])
# new manager instance reads from disk
wm2 = WorkspaceManager(root=tmp_path)
targets = wm2.list_targets(name)
assert "10.0.0.1" in targets
assert "host.local" in targets
def test_last_target_persistence(tmp_path: Path):
wm = WorkspaceManager(root=tmp_path)
a = "wsA"
b = "wsB"
wm.create(a)
wm.create(b)
t1 = "192.168.0.4"
t2 = "192.168.0.165"
# set last target on workspace A and B
norm1 = wm.set_last_target(a, t1)
norm2 = wm.set_last_target(b, t2)
# persisted in meta
assert wm.get_meta_field(a, "last_target") == norm1
assert wm.get_meta_field(b, "last_target") == norm2
# targets list contains the last target
assert norm1 in wm.list_targets(a)
assert norm2 in wm.list_targets(b)
# new manager instance still sees last_target
wm2 = WorkspaceManager(root=tmp_path)
assert wm2.get_meta_field(a, "last_target") == norm1
assert wm2.get_meta_field(b, "last_target") == norm2