axio-tui-guards 0.1.0__tar.gz
This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
- axio_tui_guards-0.1.0/.github/workflows/publish.yml +40 -0
- axio_tui_guards-0.1.0/.github/workflows/tests.yml +44 -0
- axio_tui_guards-0.1.0/LICENSE +21 -0
- axio_tui_guards-0.1.0/Makefile +16 -0
- axio_tui_guards-0.1.0/PKG-INFO +9 -0
- axio_tui_guards-0.1.0/README.md +15 -0
- axio_tui_guards-0.1.0/pyproject.toml +35 -0
- axio_tui_guards-0.1.0/src/axio_tui_guards/__init__.py +5 -0
- axio_tui_guards-0.1.0/src/axio_tui_guards/dialogs.py +160 -0
- axio_tui_guards-0.1.0/src/axio_tui_guards/guards.py +147 -0
- axio_tui_guards-0.1.0/tests/conftest.py +1 -0
- axio_tui_guards-0.1.0/tests/test_tui_guards.py +210 -0
|
@@ -0,0 +1,40 @@
|
|
|
1
|
+
name: Publish to PyPI
|
|
2
|
+
|
|
3
|
+
on:
|
|
4
|
+
release:
|
|
5
|
+
types: [published]
|
|
6
|
+
|
|
7
|
+
env:
|
|
8
|
+
FORCE_COLOR: 1
|
|
9
|
+
|
|
10
|
+
jobs:
|
|
11
|
+
build:
|
|
12
|
+
runs-on: ubuntu-latest
|
|
13
|
+
steps:
|
|
14
|
+
- uses: actions/checkout@v4
|
|
15
|
+
- uses: astral-sh/setup-uv@v6
|
|
16
|
+
|
|
17
|
+
- name: Set version from release tag
|
|
18
|
+
run: uv version "${GITHUB_REF_NAME#v}"
|
|
19
|
+
|
|
20
|
+
- name: Build package
|
|
21
|
+
run: uv build
|
|
22
|
+
|
|
23
|
+
- uses: actions/upload-artifact@v4
|
|
24
|
+
with:
|
|
25
|
+
name: dist
|
|
26
|
+
path: dist/
|
|
27
|
+
|
|
28
|
+
publish:
|
|
29
|
+
runs-on: ubuntu-latest
|
|
30
|
+
needs: build
|
|
31
|
+
environment: pypi
|
|
32
|
+
permissions:
|
|
33
|
+
id-token: write
|
|
34
|
+
steps:
|
|
35
|
+
- uses: actions/download-artifact@v4
|
|
36
|
+
with:
|
|
37
|
+
name: dist
|
|
38
|
+
path: dist/
|
|
39
|
+
|
|
40
|
+
- uses: pypa/gh-action-pypi-publish@release/v1
|
|
@@ -0,0 +1,44 @@
|
|
|
1
|
+
name: tests
|
|
2
|
+
|
|
3
|
+
on:
|
|
4
|
+
push:
|
|
5
|
+
branches: [ master, main ]
|
|
6
|
+
pull_request:
|
|
7
|
+
branches: [ master, main ]
|
|
8
|
+
|
|
9
|
+
env:
|
|
10
|
+
FORCE_COLOR: 1
|
|
11
|
+
|
|
12
|
+
jobs:
|
|
13
|
+
ruff:
|
|
14
|
+
runs-on: ubuntu-latest
|
|
15
|
+
steps:
|
|
16
|
+
- uses: actions/checkout@v4
|
|
17
|
+
- uses: astral-sh/setup-uv@v6
|
|
18
|
+
- run: uv sync --frozen
|
|
19
|
+
- run: uv run ruff check
|
|
20
|
+
- run: uv run ruff format --check
|
|
21
|
+
|
|
22
|
+
mypy:
|
|
23
|
+
runs-on: ubuntu-latest
|
|
24
|
+
steps:
|
|
25
|
+
- uses: actions/checkout@v4
|
|
26
|
+
- uses: astral-sh/setup-uv@v6
|
|
27
|
+
- run: uv sync --frozen
|
|
28
|
+
- run: uv run mypy .
|
|
29
|
+
|
|
30
|
+
tests:
|
|
31
|
+
runs-on: ubuntu-latest
|
|
32
|
+
strategy:
|
|
33
|
+
fail-fast: false
|
|
34
|
+
matrix:
|
|
35
|
+
python:
|
|
36
|
+
- "3.12"
|
|
37
|
+
- "3.13"
|
|
38
|
+
steps:
|
|
39
|
+
- uses: actions/checkout@v4
|
|
40
|
+
- uses: astral-sh/setup-uv@v6
|
|
41
|
+
with:
|
|
42
|
+
python-version: ${{ matrix.python }}
|
|
43
|
+
- run: uv sync --frozen
|
|
44
|
+
- run: uv run pytest -vv --cov=axio_tui_guards --cov-report=term-missing
|
|
@@ -0,0 +1,21 @@
|
|
|
1
|
+
MIT License
|
|
2
|
+
|
|
3
|
+
Copyright (c) 2025 Axio contributors
|
|
4
|
+
|
|
5
|
+
Permission is hereby granted, free of charge, to any person obtaining a copy
|
|
6
|
+
of this software and associated documentation files (the "Software"), to deal
|
|
7
|
+
in the Software without restriction, including without limitation the rights
|
|
8
|
+
to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
|
|
9
|
+
copies of the Software, and to permit persons to whom the Software is
|
|
10
|
+
furnished to do so, subject to the following conditions:
|
|
11
|
+
|
|
12
|
+
The above copyright notice and this permission notice shall be included in all
|
|
13
|
+
copies or substantial portions of the Software.
|
|
14
|
+
|
|
15
|
+
THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
|
|
16
|
+
IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
|
|
17
|
+
FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
|
|
18
|
+
AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
|
|
19
|
+
LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
|
|
20
|
+
OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
|
|
21
|
+
SOFTWARE.
|
|
@@ -0,0 +1,16 @@
|
|
|
1
|
+
.PHONY: fmt lint typecheck test all
|
|
2
|
+
|
|
3
|
+
fmt:
|
|
4
|
+
uv run ruff format src/ tests/
|
|
5
|
+
uv run ruff check --fix src/ tests/
|
|
6
|
+
|
|
7
|
+
lint:
|
|
8
|
+
uv run ruff check src/ tests/
|
|
9
|
+
|
|
10
|
+
typecheck:
|
|
11
|
+
uv run mypy src/
|
|
12
|
+
|
|
13
|
+
test:
|
|
14
|
+
uv run pytest tests/ -v
|
|
15
|
+
|
|
16
|
+
all: fmt lint typecheck test
|
|
@@ -0,0 +1,35 @@
|
|
|
1
|
+
[project]
|
|
2
|
+
name = "axio-tui-guards"
|
|
3
|
+
version = "0.1.0"
|
|
4
|
+
description = "Guards plugin for Axio TUI"
|
|
5
|
+
requires-python = ">=3.12"
|
|
6
|
+
license = {text = "MIT"}
|
|
7
|
+
dependencies = ["axio", "axio-tui"]
|
|
8
|
+
|
|
9
|
+
[project.entry-points."axio.guards"]
|
|
10
|
+
path = "axio_tui_guards.guards:PathGuard"
|
|
11
|
+
llm = "axio_tui_guards.guards:LLMGuard"
|
|
12
|
+
|
|
13
|
+
[build-system]
|
|
14
|
+
requires = ["hatchling"]
|
|
15
|
+
build-backend = "hatchling.build"
|
|
16
|
+
|
|
17
|
+
[tool.hatch.build.targets.wheel]
|
|
18
|
+
packages = ["src/axio_tui_guards"]
|
|
19
|
+
|
|
20
|
+
[tool.pytest.ini_options]
|
|
21
|
+
asyncio_mode = "auto"
|
|
22
|
+
|
|
23
|
+
[tool.ruff]
|
|
24
|
+
line-length = 119
|
|
25
|
+
target-version = "py312"
|
|
26
|
+
|
|
27
|
+
[tool.ruff.lint]
|
|
28
|
+
select = ["E", "F", "I", "UP"]
|
|
29
|
+
|
|
30
|
+
[tool.mypy]
|
|
31
|
+
strict = true
|
|
32
|
+
python_version = "3.12"
|
|
33
|
+
|
|
34
|
+
[dependency-groups]
|
|
35
|
+
dev = ["pytest>=8", "pytest-asyncio>=0.24", "mypy>=1.14", "ruff>=0.9"]
|
|
@@ -0,0 +1,160 @@
|
|
|
1
|
+
"""Guard dialog screens for the TUI."""
|
|
2
|
+
|
|
3
|
+
from __future__ import annotations
|
|
4
|
+
|
|
5
|
+
from textual.app import ComposeResult
|
|
6
|
+
from textual.binding import Binding
|
|
7
|
+
from textual.containers import Container, Horizontal
|
|
8
|
+
from textual.screen import ModalScreen
|
|
9
|
+
from textual.widgets import Button, Input, Static
|
|
10
|
+
|
|
11
|
+
|
|
12
|
+
class PathGuardDialog(ModalScreen[str]):
|
|
13
|
+
"""Modal dialog for path access permission — buttons only."""
|
|
14
|
+
|
|
15
|
+
BINDINGS = [
|
|
16
|
+
Binding("a", "allow", "Allow", show=False),
|
|
17
|
+
Binding("d", "deny", "Deny", show=False),
|
|
18
|
+
Binding("escape", "deny", "Deny", show=False),
|
|
19
|
+
Binding("left", "focus_prev_button", show=False),
|
|
20
|
+
Binding("right", "focus_next_button", show=False),
|
|
21
|
+
]
|
|
22
|
+
CSS = """
|
|
23
|
+
PathGuardDialog { align: center middle; }
|
|
24
|
+
#path-guard-dialog {
|
|
25
|
+
width: 80;
|
|
26
|
+
height: auto;
|
|
27
|
+
max-height: 80%;
|
|
28
|
+
border: heavy $warning;
|
|
29
|
+
background: $panel;
|
|
30
|
+
padding: 1 2;
|
|
31
|
+
}
|
|
32
|
+
.guard-buttons { height: auto; margin-top: 1; }
|
|
33
|
+
.guard-buttons Button { margin: 0 1; }
|
|
34
|
+
.guard-prompt { color: $warning; }
|
|
35
|
+
"""
|
|
36
|
+
|
|
37
|
+
def __init__(self, prompt_text: str) -> None:
|
|
38
|
+
super().__init__()
|
|
39
|
+
self._prompt_text = prompt_text
|
|
40
|
+
|
|
41
|
+
def compose(self) -> ComposeResult:
|
|
42
|
+
prompt = self._prompt_text[:5000] + "..." if len(self._prompt_text) > 5000 else self._prompt_text
|
|
43
|
+
with Container(id="path-guard-dialog"):
|
|
44
|
+
yield Static("[bold]Path Access Request[/]")
|
|
45
|
+
yield Static(prompt, markup=False, classes="guard-prompt")
|
|
46
|
+
with Horizontal(classes="guard-buttons"):
|
|
47
|
+
yield Button("Allow", id="btn-allow", variant="success")
|
|
48
|
+
yield Button("Deny", id="btn-deny", variant="error")
|
|
49
|
+
yield Button("Always Deny", id="btn-always-deny", variant="warning")
|
|
50
|
+
|
|
51
|
+
def on_mount(self) -> None:
|
|
52
|
+
self.query_one("#btn-allow", Button).focus()
|
|
53
|
+
|
|
54
|
+
def _cycle_buttons(self, direction: int) -> None:
|
|
55
|
+
buttons = list(self.query(Button))
|
|
56
|
+
if not buttons:
|
|
57
|
+
return
|
|
58
|
+
try:
|
|
59
|
+
idx = buttons.index(self.focused) # type: ignore[arg-type]
|
|
60
|
+
except ValueError:
|
|
61
|
+
buttons[0].focus()
|
|
62
|
+
return
|
|
63
|
+
buttons[(idx + direction) % len(buttons)].focus()
|
|
64
|
+
|
|
65
|
+
def action_focus_next_button(self) -> None:
|
|
66
|
+
self._cycle_buttons(1)
|
|
67
|
+
|
|
68
|
+
def action_focus_prev_button(self) -> None:
|
|
69
|
+
self._cycle_buttons(-1)
|
|
70
|
+
|
|
71
|
+
def on_button_pressed(self, event: Button.Pressed) -> None:
|
|
72
|
+
match event.button.id:
|
|
73
|
+
case "btn-allow":
|
|
74
|
+
self.dismiss("y")
|
|
75
|
+
case "btn-deny":
|
|
76
|
+
self.dismiss("n")
|
|
77
|
+
case "btn-always-deny":
|
|
78
|
+
self.dismiss("deny")
|
|
79
|
+
|
|
80
|
+
def action_allow(self) -> None:
|
|
81
|
+
self.dismiss("y")
|
|
82
|
+
|
|
83
|
+
def action_deny(self) -> None:
|
|
84
|
+
self.dismiss("n")
|
|
85
|
+
|
|
86
|
+
|
|
87
|
+
class LLMGuardDialog(ModalScreen[str]):
|
|
88
|
+
"""Modal dialog for LLM safety review — buttons + text input."""
|
|
89
|
+
|
|
90
|
+
BINDINGS = [
|
|
91
|
+
Binding("escape", "deny", "Deny", show=False),
|
|
92
|
+
Binding("left", "focus_prev_button", show=False),
|
|
93
|
+
Binding("right", "focus_next_button", show=False),
|
|
94
|
+
]
|
|
95
|
+
CSS = """
|
|
96
|
+
LLMGuardDialog { align: center middle; }
|
|
97
|
+
#llm-guard-dialog {
|
|
98
|
+
width: 80;
|
|
99
|
+
height: auto;
|
|
100
|
+
max-height: 80%;
|
|
101
|
+
border: heavy $warning;
|
|
102
|
+
background: $panel;
|
|
103
|
+
padding: 1 2;
|
|
104
|
+
}
|
|
105
|
+
.guard-buttons { height: auto; margin-top: 1; }
|
|
106
|
+
.guard-buttons Button { margin: 0 1; }
|
|
107
|
+
#guard-reason { margin-top: 1; }
|
|
108
|
+
.guard-prompt { color: $warning; }
|
|
109
|
+
"""
|
|
110
|
+
|
|
111
|
+
def __init__(self, prompt_text: str) -> None:
|
|
112
|
+
super().__init__()
|
|
113
|
+
self._prompt_text = prompt_text
|
|
114
|
+
|
|
115
|
+
def compose(self) -> ComposeResult:
|
|
116
|
+
prompt = self._prompt_text[:5000] + "..." if len(self._prompt_text) > 5000 else self._prompt_text
|
|
117
|
+
with Container(id="llm-guard-dialog"):
|
|
118
|
+
yield Static("[bold]Safety Review[/]")
|
|
119
|
+
yield Static(prompt, markup=False, classes="guard-prompt")
|
|
120
|
+
with Horizontal(classes="guard-buttons"):
|
|
121
|
+
yield Button("Allow", id="btn-allow", variant="success")
|
|
122
|
+
yield Button("Always Allow", id="btn-always", variant="primary")
|
|
123
|
+
yield Button("Deny", id="btn-deny", variant="error")
|
|
124
|
+
yield Input(placeholder="Custom reason...", id="guard-reason")
|
|
125
|
+
|
|
126
|
+
def on_mount(self) -> None:
|
|
127
|
+
self.query_one("#btn-allow", Button).focus()
|
|
128
|
+
|
|
129
|
+
def _cycle_buttons(self, direction: int) -> None:
|
|
130
|
+
buttons = list(self.query(Button))
|
|
131
|
+
if not buttons:
|
|
132
|
+
return
|
|
133
|
+
try:
|
|
134
|
+
idx = buttons.index(self.focused) # type: ignore[arg-type]
|
|
135
|
+
except ValueError:
|
|
136
|
+
buttons[0].focus()
|
|
137
|
+
return
|
|
138
|
+
buttons[(idx + direction) % len(buttons)].focus()
|
|
139
|
+
|
|
140
|
+
def action_focus_next_button(self) -> None:
|
|
141
|
+
self._cycle_buttons(1)
|
|
142
|
+
|
|
143
|
+
def action_focus_prev_button(self) -> None:
|
|
144
|
+
self._cycle_buttons(-1)
|
|
145
|
+
|
|
146
|
+
def on_button_pressed(self, event: Button.Pressed) -> None:
|
|
147
|
+
match event.button.id:
|
|
148
|
+
case "btn-allow":
|
|
149
|
+
self.dismiss("y")
|
|
150
|
+
case "btn-always":
|
|
151
|
+
self.dismiss("always")
|
|
152
|
+
case "btn-deny":
|
|
153
|
+
self.dismiss("n")
|
|
154
|
+
|
|
155
|
+
def on_input_submitted(self, message: Input.Submitted) -> None:
|
|
156
|
+
if message.value.strip():
|
|
157
|
+
self.dismiss(message.value.strip())
|
|
158
|
+
|
|
159
|
+
def action_deny(self) -> None:
|
|
160
|
+
self.dismiss("n")
|
|
@@ -0,0 +1,147 @@
|
|
|
1
|
+
"""Permission guards for the Axio TUI."""
|
|
2
|
+
|
|
3
|
+
import asyncio
|
|
4
|
+
import json
|
|
5
|
+
import threading
|
|
6
|
+
from collections.abc import Awaitable, Callable
|
|
7
|
+
from pathlib import Path
|
|
8
|
+
from typing import Any
|
|
9
|
+
|
|
10
|
+
from axio.agent import Agent
|
|
11
|
+
from axio.blocks import TextBlock, ToolUseBlock
|
|
12
|
+
from axio.context import ContextStore
|
|
13
|
+
from axio.exceptions import GuardError
|
|
14
|
+
from axio.messages import Message
|
|
15
|
+
from axio.permission import PermissionGuard
|
|
16
|
+
from axio_tui.tools import Confirm
|
|
17
|
+
|
|
18
|
+
type PromptFn = Callable[[str], Awaitable[str]]
|
|
19
|
+
|
|
20
|
+
|
|
21
|
+
class PathGuard(PermissionGuard):
|
|
22
|
+
"""Ask user about path access; allow grants access to the parent directory."""
|
|
23
|
+
|
|
24
|
+
PATH_FIELDS = ("file_path", "filename", "directory", "path", "cwd")
|
|
25
|
+
|
|
26
|
+
def __init__(self, prompt_fn: PromptFn | None = None) -> None:
|
|
27
|
+
self.allowed: set[str] = set()
|
|
28
|
+
self.denied: set[str] = set()
|
|
29
|
+
self._prompt = prompt_fn or ask_user
|
|
30
|
+
|
|
31
|
+
def _extract_path(self, handler: Any) -> tuple[str | None, str | None]:
|
|
32
|
+
data = handler.model_dump()
|
|
33
|
+
for name in self.PATH_FIELDS:
|
|
34
|
+
if name in data:
|
|
35
|
+
return str(data[name]), name
|
|
36
|
+
return None, None
|
|
37
|
+
|
|
38
|
+
@staticmethod
|
|
39
|
+
def _parent_dir(path: str) -> str:
|
|
40
|
+
p = Path(path)
|
|
41
|
+
return str(p if p.suffix == "" else p.parent)
|
|
42
|
+
|
|
43
|
+
def _is_allowed(self, path: str) -> bool:
|
|
44
|
+
d = Path(self._parent_dir(path))
|
|
45
|
+
return any(d == Path(a) or Path(a) in d.parents for a in self.allowed)
|
|
46
|
+
|
|
47
|
+
def _is_denied(self, path: str) -> bool:
|
|
48
|
+
return path in self.denied
|
|
49
|
+
|
|
50
|
+
async def check(self, handler: Any) -> Any:
|
|
51
|
+
path, field = self._extract_path(handler)
|
|
52
|
+
if path is None or self._is_allowed(path):
|
|
53
|
+
return handler
|
|
54
|
+
if self._is_denied(path):
|
|
55
|
+
raise GuardError(f"Path denied: {field}={path!r}")
|
|
56
|
+
|
|
57
|
+
directory = self._parent_dir(path)
|
|
58
|
+
msg = f"{field}={path!r} (directory: {directory})"
|
|
59
|
+
answer = (await self._prompt(msg)).strip()
|
|
60
|
+
if answer.lower() == "deny":
|
|
61
|
+
self.denied.add(path)
|
|
62
|
+
raise GuardError(f"Path denied: {field}={path!r}")
|
|
63
|
+
if not answer or answer.lower() == "n":
|
|
64
|
+
raise GuardError(f"Path access denied: {field}={path!r}")
|
|
65
|
+
# "y" or anything else — allow this directory for future calls
|
|
66
|
+
self.allowed.add(directory)
|
|
67
|
+
return handler
|
|
68
|
+
|
|
69
|
+
|
|
70
|
+
class LLMGuard(PermissionGuard):
|
|
71
|
+
"""Agent-based guard. User overrides feed back into the context for learning."""
|
|
72
|
+
|
|
73
|
+
def __init__(self, agent: Agent, context: ContextStore, prompt_fn: PromptFn | None = None) -> None:
|
|
74
|
+
self.agent = agent
|
|
75
|
+
self.context = context
|
|
76
|
+
self.allowed: set[str] = set()
|
|
77
|
+
self._prompt = prompt_fn or ask_user
|
|
78
|
+
|
|
79
|
+
async def extract_confirm(self, context: ContextStore) -> Confirm:
|
|
80
|
+
"""Find last confirm tool call from forked context."""
|
|
81
|
+
history = await context.get_history()
|
|
82
|
+
for msg in reversed(history):
|
|
83
|
+
if msg.role == "assistant":
|
|
84
|
+
for block in msg.content:
|
|
85
|
+
if isinstance(block, ToolUseBlock) and block.name == "confirm":
|
|
86
|
+
try:
|
|
87
|
+
return Confirm.model_validate(block.input)
|
|
88
|
+
except Exception:
|
|
89
|
+
return Confirm(verdict="RISKY", reason="Unparseable", category="unknown")
|
|
90
|
+
return Confirm(verdict="SAFE", reason="No verdict provided", category="unknown")
|
|
91
|
+
|
|
92
|
+
async def check(self, handler: Any) -> Any:
|
|
93
|
+
args_str = json.dumps(handler.model_dump(), default=str)
|
|
94
|
+
if len(args_str) > 2000:
|
|
95
|
+
args_str = args_str[:2000] + "..."
|
|
96
|
+
|
|
97
|
+
description = f"Tool: {type(handler).__name__}\nArguments: {args_str}"
|
|
98
|
+
if self.allowed:
|
|
99
|
+
description += "\n\nAuto-approved categories (classify as SAFE): " + ", ".join(sorted(self.allowed))
|
|
100
|
+
|
|
101
|
+
# Fork so tool-call noise is discarded; user answers persist in the original
|
|
102
|
+
forked = await self.context.fork()
|
|
103
|
+
async for _ in self.agent.run_stream(description, forked):
|
|
104
|
+
pass
|
|
105
|
+
|
|
106
|
+
confirm = await self.extract_confirm(forked)
|
|
107
|
+
|
|
108
|
+
if confirm.verdict == "SAFE":
|
|
109
|
+
return handler
|
|
110
|
+
if confirm.verdict == "DENY":
|
|
111
|
+
raise GuardError(f"DENIED: {confirm.reason}")
|
|
112
|
+
|
|
113
|
+
# RISKY — ask user
|
|
114
|
+
if confirm.category in self.allowed:
|
|
115
|
+
return handler
|
|
116
|
+
|
|
117
|
+
handler_repr = repr(handler)
|
|
118
|
+
if len(handler_repr) > 2000:
|
|
119
|
+
handler_repr = handler_repr[:2000] + "..."
|
|
120
|
+
msg = f"{handler_repr}\n\n{confirm.reason}"
|
|
121
|
+
answer = (await self._prompt(msg)).strip()
|
|
122
|
+
|
|
123
|
+
if not answer or answer.lower() == "n":
|
|
124
|
+
raise GuardError(f"User denied: {confirm.reason}")
|
|
125
|
+
|
|
126
|
+
if answer.lower() == "always":
|
|
127
|
+
self.allowed.add(confirm.category)
|
|
128
|
+
elif answer.lower() != "y":
|
|
129
|
+
# Feed user answer back and let it run more turns
|
|
130
|
+
async for _ in self.agent.run_stream(answer, forked):
|
|
131
|
+
pass
|
|
132
|
+
# Persist the user note in original context for future checks
|
|
133
|
+
await self.context.append(Message(role="user", content=[TextBlock(text=answer)]))
|
|
134
|
+
|
|
135
|
+
return handler
|
|
136
|
+
|
|
137
|
+
|
|
138
|
+
# Helper function for user input
|
|
139
|
+
ASK_LOCK = threading.Lock()
|
|
140
|
+
|
|
141
|
+
|
|
142
|
+
async def ask_user(message: str) -> str:
|
|
143
|
+
def _blocking() -> str:
|
|
144
|
+
with ASK_LOCK:
|
|
145
|
+
return input(message)
|
|
146
|
+
|
|
147
|
+
return await asyncio.to_thread(_blocking)
|
|
@@ -0,0 +1 @@
|
|
|
1
|
+
"""Shared test fixtures for axio-tui-guards."""
|
|
@@ -0,0 +1,210 @@
|
|
|
1
|
+
"""Tests for axio_tui_guards — PathGuard and LLMGuard."""
|
|
2
|
+
|
|
3
|
+
from __future__ import annotations
|
|
4
|
+
|
|
5
|
+
from typing import Any
|
|
6
|
+
|
|
7
|
+
import pytest
|
|
8
|
+
from axio.blocks import TextBlock, ToolUseBlock
|
|
9
|
+
from axio.context import MemoryContextStore
|
|
10
|
+
from axio.exceptions import GuardError
|
|
11
|
+
from axio.messages import Message
|
|
12
|
+
from axio.testing import StubTransport, make_text_response, make_tool_use_response
|
|
13
|
+
from axio_tools_local.read_file import ReadFile
|
|
14
|
+
from axio_tools_local.shell import Shell
|
|
15
|
+
from axio_tools_local.write_file import WriteFile
|
|
16
|
+
from axio_tui.tools import Confirm
|
|
17
|
+
from axio_tui_guards.guards import LLMGuard, PathGuard
|
|
18
|
+
|
|
19
|
+
|
|
20
|
+
class TestPathGuard:
|
|
21
|
+
@staticmethod
|
|
22
|
+
async def _allow(_msg: str) -> str:
|
|
23
|
+
return "y"
|
|
24
|
+
|
|
25
|
+
@staticmethod
|
|
26
|
+
async def _deny(_msg: str) -> str:
|
|
27
|
+
return "n"
|
|
28
|
+
|
|
29
|
+
@staticmethod
|
|
30
|
+
async def _always_deny(_msg: str) -> str:
|
|
31
|
+
return "deny"
|
|
32
|
+
|
|
33
|
+
async def test_allow_grants_directory(self) -> None:
|
|
34
|
+
guard = PathGuard(prompt_fn=self._allow)
|
|
35
|
+
handler = ReadFile(filename="/tmp/test/file.txt")
|
|
36
|
+
result = await guard.check(handler)
|
|
37
|
+
assert result is handler
|
|
38
|
+
# Same directory should be auto-allowed now
|
|
39
|
+
handler2 = ReadFile(filename="/tmp/test/other.txt")
|
|
40
|
+
result2 = await guard.check(handler2)
|
|
41
|
+
assert result2 is handler2
|
|
42
|
+
|
|
43
|
+
async def test_deny_raises_guard_error(self) -> None:
|
|
44
|
+
guard = PathGuard(prompt_fn=self._deny)
|
|
45
|
+
handler = ReadFile(filename="/secret/file.txt")
|
|
46
|
+
with pytest.raises(GuardError, match="denied"):
|
|
47
|
+
await guard.check(handler)
|
|
48
|
+
|
|
49
|
+
async def test_always_deny_persists(self) -> None:
|
|
50
|
+
guard = PathGuard(prompt_fn=self._always_deny)
|
|
51
|
+
handler = ReadFile(filename="/deny/file.txt")
|
|
52
|
+
with pytest.raises(GuardError):
|
|
53
|
+
await guard.check(handler)
|
|
54
|
+
# Second call should also be denied without prompting
|
|
55
|
+
with pytest.raises(GuardError, match="denied"):
|
|
56
|
+
await guard.check(handler)
|
|
57
|
+
|
|
58
|
+
async def test_no_path_field_passes_through(self) -> None:
|
|
59
|
+
guard = PathGuard(prompt_fn=self._deny)
|
|
60
|
+
handler = Confirm(verdict="SAFE", reason="ok", category="test")
|
|
61
|
+
result = await guard.check(handler)
|
|
62
|
+
assert result is handler
|
|
63
|
+
|
|
64
|
+
async def test_shell_extracts_cwd(self) -> None:
|
|
65
|
+
guard = PathGuard(prompt_fn=self._allow)
|
|
66
|
+
handler = Shell(command="ls", cwd="/tmp/project")
|
|
67
|
+
result = await guard.check(handler)
|
|
68
|
+
assert result is handler
|
|
69
|
+
assert "/tmp/project" in guard.allowed
|
|
70
|
+
|
|
71
|
+
async def test_subdirectory_auto_allowed(self) -> None:
|
|
72
|
+
guard = PathGuard(prompt_fn=self._allow)
|
|
73
|
+
handler = WriteFile(file_path="/home/user/project/src/main.py", content="x")
|
|
74
|
+
await guard.check(handler)
|
|
75
|
+
# Child path in same tree should be auto-allowed
|
|
76
|
+
handler2 = WriteFile(file_path="/home/user/project/src/lib/util.py", content="y")
|
|
77
|
+
result = await guard.check(handler2)
|
|
78
|
+
assert result is handler2
|
|
79
|
+
|
|
80
|
+
|
|
81
|
+
class TestLLMGuard:
|
|
82
|
+
async def test_safe_verdict_passes(self) -> None:
|
|
83
|
+
confirm_input = {"verdict": "SAFE", "reason": "harmless", "category": "read"}
|
|
84
|
+
transport = StubTransport(
|
|
85
|
+
[
|
|
86
|
+
make_tool_use_response("confirm", "call_1", confirm_input),
|
|
87
|
+
make_text_response("ok"),
|
|
88
|
+
]
|
|
89
|
+
)
|
|
90
|
+
agent_from_transport = _make_guard_agent(transport)
|
|
91
|
+
guard = LLMGuard(agent_from_transport, MemoryContextStore())
|
|
92
|
+
handler = ReadFile(filename="test.txt")
|
|
93
|
+
result = await guard.check(handler)
|
|
94
|
+
assert result is handler
|
|
95
|
+
|
|
96
|
+
async def test_deny_verdict_raises(self) -> None:
|
|
97
|
+
confirm_input = {"verdict": "DENY", "reason": "malicious", "category": "exec"}
|
|
98
|
+
transport = StubTransport(
|
|
99
|
+
[
|
|
100
|
+
make_tool_use_response("confirm", "call_1", confirm_input),
|
|
101
|
+
make_text_response("blocked"),
|
|
102
|
+
]
|
|
103
|
+
)
|
|
104
|
+
agent = _make_guard_agent(transport)
|
|
105
|
+
guard = LLMGuard(agent, MemoryContextStore())
|
|
106
|
+
handler = Shell(command="rm -rf /")
|
|
107
|
+
with pytest.raises(GuardError, match="DENIED"):
|
|
108
|
+
await guard.check(handler)
|
|
109
|
+
|
|
110
|
+
async def test_risky_verdict_user_allows(self) -> None:
|
|
111
|
+
confirm_input = {"verdict": "RISKY", "reason": "writes file", "category": "write"}
|
|
112
|
+
transport = StubTransport(
|
|
113
|
+
[
|
|
114
|
+
make_tool_use_response("confirm", "call_1", confirm_input),
|
|
115
|
+
make_text_response("review needed"),
|
|
116
|
+
]
|
|
117
|
+
)
|
|
118
|
+
agent = _make_guard_agent(transport)
|
|
119
|
+
|
|
120
|
+
async def allow(_msg: str) -> str:
|
|
121
|
+
return "y"
|
|
122
|
+
|
|
123
|
+
guard = LLMGuard(agent, MemoryContextStore(), prompt_fn=allow)
|
|
124
|
+
handler = WriteFile(file_path="out.txt", content="data")
|
|
125
|
+
result = await guard.check(handler)
|
|
126
|
+
assert result is handler
|
|
127
|
+
|
|
128
|
+
async def test_risky_verdict_user_denies(self) -> None:
|
|
129
|
+
confirm_input = {"verdict": "RISKY", "reason": "dangerous", "category": "exec"}
|
|
130
|
+
transport = StubTransport(
|
|
131
|
+
[
|
|
132
|
+
make_tool_use_response("confirm", "call_1", confirm_input),
|
|
133
|
+
make_text_response("needs review"),
|
|
134
|
+
]
|
|
135
|
+
)
|
|
136
|
+
agent = _make_guard_agent(transport)
|
|
137
|
+
|
|
138
|
+
async def deny(_msg: str) -> str:
|
|
139
|
+
return "n"
|
|
140
|
+
|
|
141
|
+
guard = LLMGuard(agent, MemoryContextStore(), prompt_fn=deny)
|
|
142
|
+
handler = Shell(command="sudo reboot")
|
|
143
|
+
with pytest.raises(GuardError, match="denied"):
|
|
144
|
+
await guard.check(handler)
|
|
145
|
+
|
|
146
|
+
async def test_always_allows_category(self) -> None:
|
|
147
|
+
confirm_input = {"verdict": "RISKY", "reason": "writes", "category": "write_file"}
|
|
148
|
+
transport = StubTransport(
|
|
149
|
+
[
|
|
150
|
+
make_tool_use_response("confirm", "call_1", confirm_input),
|
|
151
|
+
make_text_response("ok"),
|
|
152
|
+
# Second call — should not reach transport since category is pre-approved
|
|
153
|
+
]
|
|
154
|
+
)
|
|
155
|
+
agent = _make_guard_agent(transport)
|
|
156
|
+
|
|
157
|
+
async def always(_msg: str) -> str:
|
|
158
|
+
return "always"
|
|
159
|
+
|
|
160
|
+
guard = LLMGuard(agent, MemoryContextStore(), prompt_fn=always)
|
|
161
|
+
handler = WriteFile(file_path="a.txt", content="data")
|
|
162
|
+
await guard.check(handler)
|
|
163
|
+
assert "write_file" in guard.allowed
|
|
164
|
+
|
|
165
|
+
async def test_extract_confirm_from_context(self) -> None:
|
|
166
|
+
guard = LLMGuard.__new__(LLMGuard)
|
|
167
|
+
ctx = MemoryContextStore()
|
|
168
|
+
await ctx.append(Message(role="user", content=[TextBlock(text="check this")]))
|
|
169
|
+
await ctx.append(
|
|
170
|
+
Message(
|
|
171
|
+
role="assistant",
|
|
172
|
+
content=[
|
|
173
|
+
ToolUseBlock(
|
|
174
|
+
id="c1", name="confirm", input={"verdict": "SAFE", "reason": "ok", "category": "read"}
|
|
175
|
+
),
|
|
176
|
+
],
|
|
177
|
+
)
|
|
178
|
+
)
|
|
179
|
+
confirm = await guard.extract_confirm(ctx)
|
|
180
|
+
assert confirm.verdict == "SAFE"
|
|
181
|
+
assert confirm.category == "read"
|
|
182
|
+
|
|
183
|
+
async def test_extract_confirm_no_verdict(self) -> None:
|
|
184
|
+
guard = LLMGuard.__new__(LLMGuard)
|
|
185
|
+
ctx = MemoryContextStore()
|
|
186
|
+
await ctx.append(Message(role="user", content=[TextBlock(text="check")]))
|
|
187
|
+
await ctx.append(Message(role="assistant", content=[TextBlock(text="no tool call")]))
|
|
188
|
+
confirm = await guard.extract_confirm(ctx)
|
|
189
|
+
assert confirm.verdict == "SAFE"
|
|
190
|
+
assert confirm.reason == "No verdict provided"
|
|
191
|
+
|
|
192
|
+
|
|
193
|
+
def _make_guard_agent(transport: Any) -> Any:
|
|
194
|
+
from axio.agent import Agent
|
|
195
|
+
from axio.tool import Tool
|
|
196
|
+
from axio_tools_local.list_files import ListFiles
|
|
197
|
+
from axio_tools_local.read_file import ReadFile
|
|
198
|
+
from axio_tui.tools import Confirm, StatusLine
|
|
199
|
+
|
|
200
|
+
return Agent(
|
|
201
|
+
system="You are a safety classifier.",
|
|
202
|
+
tools=[
|
|
203
|
+
Tool(name="status_line", description="Set status", handler=StatusLine),
|
|
204
|
+
Tool(name="read_file", description="Read file", handler=ReadFile),
|
|
205
|
+
Tool(name="list", description="List files", handler=ListFiles),
|
|
206
|
+
Tool(name="confirm", description="Submit verdict", handler=Confirm),
|
|
207
|
+
],
|
|
208
|
+
transport=transport,
|
|
209
|
+
max_iterations=5,
|
|
210
|
+
)
|