aiptx 2.0.7__py3-none-any.whl

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.
Files changed (187) hide show
  1. aipt_v2/__init__.py +110 -0
  2. aipt_v2/__main__.py +24 -0
  3. aipt_v2/agents/AIPTxAgent/__init__.py +10 -0
  4. aipt_v2/agents/AIPTxAgent/aiptx_agent.py +211 -0
  5. aipt_v2/agents/__init__.py +46 -0
  6. aipt_v2/agents/base.py +520 -0
  7. aipt_v2/agents/exploit_agent.py +688 -0
  8. aipt_v2/agents/ptt.py +406 -0
  9. aipt_v2/agents/state.py +168 -0
  10. aipt_v2/app.py +957 -0
  11. aipt_v2/browser/__init__.py +31 -0
  12. aipt_v2/browser/automation.py +458 -0
  13. aipt_v2/browser/crawler.py +453 -0
  14. aipt_v2/cli.py +2933 -0
  15. aipt_v2/compliance/__init__.py +71 -0
  16. aipt_v2/compliance/compliance_report.py +449 -0
  17. aipt_v2/compliance/framework_mapper.py +424 -0
  18. aipt_v2/compliance/nist_mapping.py +345 -0
  19. aipt_v2/compliance/owasp_mapping.py +330 -0
  20. aipt_v2/compliance/pci_mapping.py +297 -0
  21. aipt_v2/config.py +341 -0
  22. aipt_v2/core/__init__.py +43 -0
  23. aipt_v2/core/agent.py +630 -0
  24. aipt_v2/core/llm.py +395 -0
  25. aipt_v2/core/memory.py +305 -0
  26. aipt_v2/core/ptt.py +329 -0
  27. aipt_v2/database/__init__.py +14 -0
  28. aipt_v2/database/models.py +232 -0
  29. aipt_v2/database/repository.py +384 -0
  30. aipt_v2/docker/__init__.py +23 -0
  31. aipt_v2/docker/builder.py +260 -0
  32. aipt_v2/docker/manager.py +222 -0
  33. aipt_v2/docker/sandbox.py +371 -0
  34. aipt_v2/evasion/__init__.py +58 -0
  35. aipt_v2/evasion/request_obfuscator.py +272 -0
  36. aipt_v2/evasion/tls_fingerprint.py +285 -0
  37. aipt_v2/evasion/ua_rotator.py +301 -0
  38. aipt_v2/evasion/waf_bypass.py +439 -0
  39. aipt_v2/execution/__init__.py +23 -0
  40. aipt_v2/execution/executor.py +302 -0
  41. aipt_v2/execution/parser.py +544 -0
  42. aipt_v2/execution/terminal.py +337 -0
  43. aipt_v2/health.py +437 -0
  44. aipt_v2/intelligence/__init__.py +194 -0
  45. aipt_v2/intelligence/adaptation.py +474 -0
  46. aipt_v2/intelligence/auth.py +520 -0
  47. aipt_v2/intelligence/chaining.py +775 -0
  48. aipt_v2/intelligence/correlation.py +536 -0
  49. aipt_v2/intelligence/cve_aipt.py +334 -0
  50. aipt_v2/intelligence/cve_info.py +1111 -0
  51. aipt_v2/intelligence/knowledge_graph.py +590 -0
  52. aipt_v2/intelligence/learning.py +626 -0
  53. aipt_v2/intelligence/llm_analyzer.py +502 -0
  54. aipt_v2/intelligence/llm_tool_selector.py +518 -0
  55. aipt_v2/intelligence/payload_generator.py +562 -0
  56. aipt_v2/intelligence/rag.py +239 -0
  57. aipt_v2/intelligence/scope.py +442 -0
  58. aipt_v2/intelligence/searchers/__init__.py +5 -0
  59. aipt_v2/intelligence/searchers/exploitdb_searcher.py +523 -0
  60. aipt_v2/intelligence/searchers/github_searcher.py +467 -0
  61. aipt_v2/intelligence/searchers/google_searcher.py +281 -0
  62. aipt_v2/intelligence/tools.json +443 -0
  63. aipt_v2/intelligence/triage.py +670 -0
  64. aipt_v2/interactive_shell.py +559 -0
  65. aipt_v2/interface/__init__.py +5 -0
  66. aipt_v2/interface/cli.py +230 -0
  67. aipt_v2/interface/main.py +501 -0
  68. aipt_v2/interface/tui.py +1276 -0
  69. aipt_v2/interface/utils.py +583 -0
  70. aipt_v2/llm/__init__.py +39 -0
  71. aipt_v2/llm/config.py +26 -0
  72. aipt_v2/llm/llm.py +514 -0
  73. aipt_v2/llm/memory.py +214 -0
  74. aipt_v2/llm/request_queue.py +89 -0
  75. aipt_v2/llm/utils.py +89 -0
  76. aipt_v2/local_tool_installer.py +1467 -0
  77. aipt_v2/models/__init__.py +15 -0
  78. aipt_v2/models/findings.py +295 -0
  79. aipt_v2/models/phase_result.py +224 -0
  80. aipt_v2/models/scan_config.py +207 -0
  81. aipt_v2/monitoring/grafana/dashboards/aipt-dashboard.json +355 -0
  82. aipt_v2/monitoring/grafana/dashboards/default.yml +17 -0
  83. aipt_v2/monitoring/grafana/datasources/prometheus.yml +17 -0
  84. aipt_v2/monitoring/prometheus.yml +60 -0
  85. aipt_v2/orchestration/__init__.py +52 -0
  86. aipt_v2/orchestration/pipeline.py +398 -0
  87. aipt_v2/orchestration/progress.py +300 -0
  88. aipt_v2/orchestration/scheduler.py +296 -0
  89. aipt_v2/orchestrator.py +2427 -0
  90. aipt_v2/payloads/__init__.py +27 -0
  91. aipt_v2/payloads/cmdi.py +150 -0
  92. aipt_v2/payloads/sqli.py +263 -0
  93. aipt_v2/payloads/ssrf.py +204 -0
  94. aipt_v2/payloads/templates.py +222 -0
  95. aipt_v2/payloads/traversal.py +166 -0
  96. aipt_v2/payloads/xss.py +204 -0
  97. aipt_v2/prompts/__init__.py +60 -0
  98. aipt_v2/proxy/__init__.py +29 -0
  99. aipt_v2/proxy/history.py +352 -0
  100. aipt_v2/proxy/interceptor.py +452 -0
  101. aipt_v2/recon/__init__.py +44 -0
  102. aipt_v2/recon/dns.py +241 -0
  103. aipt_v2/recon/osint.py +367 -0
  104. aipt_v2/recon/subdomain.py +372 -0
  105. aipt_v2/recon/tech_detect.py +311 -0
  106. aipt_v2/reports/__init__.py +17 -0
  107. aipt_v2/reports/generator.py +313 -0
  108. aipt_v2/reports/html_report.py +378 -0
  109. aipt_v2/runtime/__init__.py +53 -0
  110. aipt_v2/runtime/base.py +30 -0
  111. aipt_v2/runtime/docker.py +401 -0
  112. aipt_v2/runtime/local.py +346 -0
  113. aipt_v2/runtime/tool_server.py +205 -0
  114. aipt_v2/runtime/vps.py +830 -0
  115. aipt_v2/scanners/__init__.py +28 -0
  116. aipt_v2/scanners/base.py +273 -0
  117. aipt_v2/scanners/nikto.py +244 -0
  118. aipt_v2/scanners/nmap.py +402 -0
  119. aipt_v2/scanners/nuclei.py +273 -0
  120. aipt_v2/scanners/web.py +454 -0
  121. aipt_v2/scripts/security_audit.py +366 -0
  122. aipt_v2/setup_wizard.py +941 -0
  123. aipt_v2/skills/__init__.py +80 -0
  124. aipt_v2/skills/agents/__init__.py +14 -0
  125. aipt_v2/skills/agents/api_tester.py +706 -0
  126. aipt_v2/skills/agents/base.py +477 -0
  127. aipt_v2/skills/agents/code_review.py +459 -0
  128. aipt_v2/skills/agents/security_agent.py +336 -0
  129. aipt_v2/skills/agents/web_pentest.py +818 -0
  130. aipt_v2/skills/prompts/__init__.py +647 -0
  131. aipt_v2/system_detector.py +539 -0
  132. aipt_v2/telemetry/__init__.py +7 -0
  133. aipt_v2/telemetry/tracer.py +347 -0
  134. aipt_v2/terminal/__init__.py +28 -0
  135. aipt_v2/terminal/executor.py +400 -0
  136. aipt_v2/terminal/sandbox.py +350 -0
  137. aipt_v2/tools/__init__.py +44 -0
  138. aipt_v2/tools/active_directory/__init__.py +78 -0
  139. aipt_v2/tools/active_directory/ad_config.py +238 -0
  140. aipt_v2/tools/active_directory/bloodhound_wrapper.py +447 -0
  141. aipt_v2/tools/active_directory/kerberos_attacks.py +430 -0
  142. aipt_v2/tools/active_directory/ldap_enum.py +533 -0
  143. aipt_v2/tools/active_directory/smb_attacks.py +505 -0
  144. aipt_v2/tools/agents_graph/__init__.py +19 -0
  145. aipt_v2/tools/agents_graph/agents_graph_actions.py +69 -0
  146. aipt_v2/tools/api_security/__init__.py +76 -0
  147. aipt_v2/tools/api_security/api_discovery.py +608 -0
  148. aipt_v2/tools/api_security/graphql_scanner.py +622 -0
  149. aipt_v2/tools/api_security/jwt_analyzer.py +577 -0
  150. aipt_v2/tools/api_security/openapi_fuzzer.py +761 -0
  151. aipt_v2/tools/browser/__init__.py +5 -0
  152. aipt_v2/tools/browser/browser_actions.py +238 -0
  153. aipt_v2/tools/browser/browser_instance.py +535 -0
  154. aipt_v2/tools/browser/tab_manager.py +344 -0
  155. aipt_v2/tools/cloud/__init__.py +70 -0
  156. aipt_v2/tools/cloud/cloud_config.py +273 -0
  157. aipt_v2/tools/cloud/cloud_scanner.py +639 -0
  158. aipt_v2/tools/cloud/prowler_tool.py +571 -0
  159. aipt_v2/tools/cloud/scoutsuite_tool.py +359 -0
  160. aipt_v2/tools/executor.py +307 -0
  161. aipt_v2/tools/parser.py +408 -0
  162. aipt_v2/tools/proxy/__init__.py +5 -0
  163. aipt_v2/tools/proxy/proxy_actions.py +103 -0
  164. aipt_v2/tools/proxy/proxy_manager.py +789 -0
  165. aipt_v2/tools/registry.py +196 -0
  166. aipt_v2/tools/scanners/__init__.py +343 -0
  167. aipt_v2/tools/scanners/acunetix_tool.py +712 -0
  168. aipt_v2/tools/scanners/burp_tool.py +631 -0
  169. aipt_v2/tools/scanners/config.py +156 -0
  170. aipt_v2/tools/scanners/nessus_tool.py +588 -0
  171. aipt_v2/tools/scanners/zap_tool.py +612 -0
  172. aipt_v2/tools/terminal/__init__.py +5 -0
  173. aipt_v2/tools/terminal/terminal_actions.py +37 -0
  174. aipt_v2/tools/terminal/terminal_manager.py +153 -0
  175. aipt_v2/tools/terminal/terminal_session.py +449 -0
  176. aipt_v2/tools/tool_processing.py +108 -0
  177. aipt_v2/utils/__init__.py +17 -0
  178. aipt_v2/utils/logging.py +202 -0
  179. aipt_v2/utils/model_manager.py +187 -0
  180. aipt_v2/utils/searchers/__init__.py +269 -0
  181. aipt_v2/verify_install.py +793 -0
  182. aiptx-2.0.7.dist-info/METADATA +345 -0
  183. aiptx-2.0.7.dist-info/RECORD +187 -0
  184. aiptx-2.0.7.dist-info/WHEEL +5 -0
  185. aiptx-2.0.7.dist-info/entry_points.txt +7 -0
  186. aiptx-2.0.7.dist-info/licenses/LICENSE +21 -0
  187. aiptx-2.0.7.dist-info/top_level.txt +1 -0
@@ -0,0 +1,1276 @@
1
+ from __future__ import annotations
2
+
3
+ import argparse
4
+ import asyncio
5
+ import atexit
6
+ import logging
7
+ import random
8
+ import signal
9
+ import sys
10
+ import threading
11
+ from collections.abc import Callable
12
+ from importlib.metadata import PackageNotFoundError
13
+ from importlib.metadata import version as pkg_version
14
+ from typing import TYPE_CHECKING, Any, ClassVar, cast
15
+
16
+
17
+ if TYPE_CHECKING:
18
+ from textual.timer import Timer
19
+
20
+ from rich.align import Align
21
+ from rich.console import Group
22
+ from rich.markup import escape as rich_escape
23
+ from rich.panel import Panel
24
+ from rich.style import Style
25
+ from rich.text import Text
26
+ from textual import events, on
27
+ from textual.app import App, ComposeResult
28
+ from textual.binding import Binding
29
+ from textual.containers import Grid, Horizontal, Vertical, VerticalScroll
30
+ from textual.reactive import reactive
31
+ from textual.screen import ModalScreen
32
+ from textual.widgets import Button, Label, Static, TextArea, Tree
33
+ from textual.widgets.tree import TreeNode
34
+
35
+ from aipt_v2.agents.AIPTxAgent import AIPTxAgent
36
+ from aipt_v2.interface.utils import build_live_stats_text
37
+ from aipt_v2.llm.config import LLMConfig
38
+ from aipt_v2.telemetry.tracer import Tracer, set_global_tracer
39
+
40
+
41
+ def escape_markup(text: str) -> str:
42
+ return cast("str", rich_escape(text))
43
+
44
+
45
+ def get_package_version() -> str:
46
+ try:
47
+ return pkg_version("aipt-agent")
48
+ except PackageNotFoundError:
49
+ return "dev"
50
+
51
+
52
+ class ChatTextArea(TextArea): # type: ignore[misc]
53
+ def __init__(self, *args: Any, **kwargs: Any) -> None:
54
+ super().__init__(*args, **kwargs)
55
+ self._app_reference: AIPTxTUIApp | None = None
56
+
57
+ def set_app_reference(self, app: "AIPTxTUIApp") -> None:
58
+ self._app_reference = app
59
+
60
+ def _on_key(self, event: events.Key) -> None:
61
+ if event.key == "enter" and self._app_reference:
62
+ text_content = str(self.text) # type: ignore[has-type]
63
+ message = text_content.strip()
64
+ if message:
65
+ self.text = ""
66
+
67
+ self._app_reference._send_user_message(message)
68
+
69
+ event.prevent_default()
70
+ return
71
+
72
+ super()._on_key(event)
73
+
74
+
75
+ class SplashScreen(Static): # type: ignore[misc]
76
+ PRIMARY_GREEN = "#22c55e"
77
+ BANNER = (
78
+ " ███████╗████████╗██████╗ ██╗██╗ ██╗\n"
79
+ " ██╔════╝╚══██╔══╝██╔══██╗██║╚██╗██╔╝\n"
80
+ " ███████╗ ██║ ██████╔╝██║ ╚███╔╝\n"
81
+ " ╚════██║ ██║ ██╔══██╗██║ ██╔██╗\n"
82
+ " ███████║ ██║ ██║ ██║██║██╔╝ ██╗\n"
83
+ " ╚══════╝ ╚═╝ ╚═╝ ╚═╝╚═╝╚═╝ ╚═╝"
84
+ )
85
+
86
+ def __init__(self, *args: Any, **kwargs: Any) -> None:
87
+ super().__init__(*args, **kwargs)
88
+ self._animation_step = 0
89
+ self._animation_timer: Timer | None = None
90
+ self._panel_static: Static | None = None
91
+ self._version = "dev"
92
+
93
+ def compose(self) -> ComposeResult:
94
+ self._version = get_package_version()
95
+ self._animation_step = 0
96
+ start_line = self._build_start_line_text(self._animation_step)
97
+ panel = self._build_panel(start_line)
98
+
99
+ panel_static = Static(panel, id="splash_content")
100
+ self._panel_static = panel_static
101
+ yield panel_static
102
+
103
+ def on_mount(self) -> None:
104
+ self._animation_timer = self.set_interval(0.45, self._animate_start_line)
105
+
106
+ def on_unmount(self) -> None:
107
+ if self._animation_timer is not None:
108
+ self._animation_timer.stop()
109
+ self._animation_timer = None
110
+
111
+ def _animate_start_line(self) -> None:
112
+ if not self._panel_static:
113
+ return
114
+
115
+ self._animation_step += 1
116
+ start_line = self._build_start_line_text(self._animation_step)
117
+ panel = self._build_panel(start_line)
118
+ self._panel_static.update(panel)
119
+
120
+ def _build_panel(self, start_line: Text) -> Panel:
121
+ content = Group(
122
+ Align.center(Text(self.BANNER.strip("\n"), style=self.PRIMARY_GREEN, justify="center")),
123
+ Align.center(Text(" ")),
124
+ Align.center(self._build_welcome_text()),
125
+ Align.center(self._build_version_text()),
126
+ Align.center(self._build_tagline_text()),
127
+ Align.center(Text(" ")),
128
+ Align.center(start_line.copy()),
129
+ )
130
+
131
+ return Panel.fit(content, border_style=self.PRIMARY_GREEN, padding=(1, 6))
132
+
133
+ def _build_welcome_text(self) -> Text:
134
+ text = Text("Welcome to ", style=Style(color="white", bold=True))
135
+ text.append("AIPT", style=Style(color=self.PRIMARY_GREEN, bold=True))
136
+ text.append("!", style=Style(color="white", bold=True))
137
+ return text
138
+
139
+ def _build_version_text(self) -> Text:
140
+ return Text(f"v{self._version}", style=Style(color="white", dim=True))
141
+
142
+ def _build_tagline_text(self) -> Text:
143
+ return Text("Open-source AI hackers for your apps", style=Style(color="white", dim=True))
144
+
145
+ def _build_start_line_text(self, phase: int) -> Text:
146
+ emphasize = phase % 2 == 1
147
+ base_style = Style(color="white", dim=not emphasize, bold=emphasize)
148
+ aiptx_style = Style(color=self.PRIMARY_GREEN, bold=bool(emphasize))
149
+
150
+ text = Text("Starting ", style=base_style)
151
+ text.append("AIPT", style=aiptx_style)
152
+ text.append(" Cybersecurity Agent", style=base_style)
153
+
154
+ return text
155
+
156
+
157
+ class HelpScreen(ModalScreen): # type: ignore[misc]
158
+ def compose(self) -> ComposeResult:
159
+ yield Grid(
160
+ Label("🔒 AIPTx Help", id="help_title"),
161
+ Label(
162
+ "F1 Help\nCtrl+Q/C Quit\nESC Stop Agent\n"
163
+ "Enter Send message to agent\nTab Switch panels\n↑/↓ Navigate tree",
164
+ id="help_content",
165
+ ),
166
+ id="dialog",
167
+ )
168
+
169
+ def on_key(self, _event: events.Key) -> None:
170
+ self.app.pop_screen()
171
+
172
+
173
+ class StopAgentScreen(ModalScreen): # type: ignore[misc]
174
+ def __init__(self, agent_name: str, agent_id: str):
175
+ super().__init__()
176
+ self.agent_name = agent_name
177
+ self.agent_id = agent_id
178
+
179
+ def compose(self) -> ComposeResult:
180
+ yield Grid(
181
+ Label(f"🛑 Stop '{self.agent_name}'?", id="stop_agent_title"),
182
+ Grid(
183
+ Button("Yes", variant="error", id="stop_agent"),
184
+ Button("No", variant="default", id="cancel_stop"),
185
+ id="stop_agent_buttons",
186
+ ),
187
+ id="stop_agent_dialog",
188
+ )
189
+
190
+ def on_mount(self) -> None:
191
+ cancel_button = self.query_one("#cancel_stop", Button)
192
+ cancel_button.focus()
193
+
194
+ def on_key(self, event: events.Key) -> None:
195
+ if event.key in ("left", "right", "up", "down"):
196
+ focused = self.focused
197
+
198
+ if focused and focused.id == "stop_agent":
199
+ cancel_button = self.query_one("#cancel_stop", Button)
200
+ cancel_button.focus()
201
+ else:
202
+ stop_button = self.query_one("#stop_agent", Button)
203
+ stop_button.focus()
204
+
205
+ event.prevent_default()
206
+ elif event.key == "enter":
207
+ focused = self.focused
208
+ if focused and isinstance(focused, Button):
209
+ focused.press()
210
+ event.prevent_default()
211
+ elif event.key == "escape":
212
+ self.app.pop_screen()
213
+ event.prevent_default()
214
+
215
+ def on_button_pressed(self, event: Button.Pressed) -> None:
216
+ if event.button.id == "stop_agent":
217
+ self.app.action_confirm_stop_agent(self.agent_id)
218
+ else:
219
+ self.app.pop_screen()
220
+
221
+
222
+ class QuitScreen(ModalScreen): # type: ignore[misc]
223
+ def compose(self) -> ComposeResult:
224
+ yield Grid(
225
+ Label("🔒 Quit AIPTx? ", id="quit_title"),
226
+ Grid(
227
+ Button("Yes", variant="error", id="quit"),
228
+ Button("No", variant="default", id="cancel"),
229
+ id="quit_buttons",
230
+ ),
231
+ id="quit_dialog",
232
+ )
233
+
234
+ def on_mount(self) -> None:
235
+ cancel_button = self.query_one("#cancel", Button)
236
+ cancel_button.focus()
237
+
238
+ def on_key(self, event: events.Key) -> None:
239
+ if event.key in ("left", "right", "up", "down"):
240
+ focused = self.focused
241
+
242
+ if focused and focused.id == "quit":
243
+ cancel_button = self.query_one("#cancel", Button)
244
+ cancel_button.focus()
245
+ else:
246
+ quit_button = self.query_one("#quit", Button)
247
+ quit_button.focus()
248
+
249
+ event.prevent_default()
250
+ elif event.key == "enter":
251
+ focused = self.focused
252
+ if focused and isinstance(focused, Button):
253
+ focused.press()
254
+ event.prevent_default()
255
+ elif event.key == "escape":
256
+ self.app.pop_screen()
257
+ event.prevent_default()
258
+
259
+ def on_button_pressed(self, event: Button.Pressed) -> None:
260
+ if event.button.id == "quit":
261
+ self.app.action_custom_quit()
262
+ else:
263
+ self.app.pop_screen()
264
+
265
+
266
+ class AIPTxTUIApp(App): # type: ignore[misc]
267
+ CSS_PATH = "assets/tui_styles.tcss"
268
+
269
+ selected_agent_id: reactive[str | None] = reactive(default=None)
270
+ show_splash: reactive[bool] = reactive(default=True)
271
+
272
+ BINDINGS: ClassVar[list[Binding]] = [
273
+ Binding("f1", "toggle_help", "Help", priority=True),
274
+ Binding("ctrl+q", "request_quit", "Quit", priority=True),
275
+ Binding("ctrl+c", "request_quit", "Quit", priority=True),
276
+ Binding("escape", "stop_selected_agent", "Stop Agent", priority=True),
277
+ ]
278
+
279
+ def __init__(self, args: argparse.Namespace):
280
+ super().__init__()
281
+ self.args = args
282
+ self.scan_config = self._build_scan_config(args)
283
+ self.agent_config = self._build_agent_config(args)
284
+
285
+ self.tracer = Tracer(self.scan_config["run_name"])
286
+ self.tracer.set_scan_config(self.scan_config)
287
+ set_global_tracer(self.tracer)
288
+
289
+ self.agent_nodes: dict[str, TreeNode] = {}
290
+
291
+ self._displayed_agents: set[str] = set()
292
+ self._displayed_events: list[str] = []
293
+
294
+ self._scan_thread: threading.Thread | None = None
295
+ self._scan_stop_event = threading.Event()
296
+ self._scan_completed = threading.Event()
297
+
298
+ self._action_verbs = [
299
+ "Generating",
300
+ "Scanning",
301
+ "Analyzing",
302
+ "Probing",
303
+ "Hacking",
304
+ "Testing",
305
+ "Exploiting",
306
+ "Investigating",
307
+ ]
308
+ self._agent_verbs: dict[str, str] = {} # agent_id -> current_verb
309
+ self._agent_verb_timers: dict[str, Any] = {} # agent_id -> timer
310
+ self._agent_dot_states: dict[str, int] = {} # agent_id -> dot_count (0-3)
311
+ self._dot_animation_timer: Any | None = None
312
+
313
+ self._setup_cleanup_handlers()
314
+
315
+ def _build_scan_config(self, args: argparse.Namespace) -> dict[str, Any]:
316
+ return {
317
+ "scan_id": args.run_name,
318
+ "targets": args.targets_info,
319
+ "user_instructions": args.instruction or "",
320
+ "run_name": args.run_name,
321
+ }
322
+
323
+ def _build_agent_config(self, args: argparse.Namespace) -> dict[str, Any]:
324
+ llm_config = LLMConfig()
325
+
326
+ config = {
327
+ "llm_config": llm_config,
328
+ "max_iterations": 300,
329
+ }
330
+
331
+ if getattr(args, "local_sources", None):
332
+ config["local_sources"] = args.local_sources
333
+
334
+ return config
335
+
336
+ def _setup_cleanup_handlers(self) -> None:
337
+ def cleanup_on_exit() -> None:
338
+ self.tracer.cleanup()
339
+
340
+ def signal_handler(_signum: int, _frame: Any) -> None:
341
+ self.tracer.cleanup()
342
+ sys.exit(0)
343
+
344
+ atexit.register(cleanup_on_exit)
345
+ signal.signal(signal.SIGINT, signal_handler)
346
+ signal.signal(signal.SIGTERM, signal_handler)
347
+ if hasattr(signal, "SIGHUP"):
348
+ signal.signal(signal.SIGHUP, signal_handler)
349
+
350
+ def compose(self) -> ComposeResult:
351
+ if self.show_splash:
352
+ yield SplashScreen(id="splash_screen")
353
+
354
+ def watch_show_splash(self, show_splash: bool) -> None:
355
+ if not show_splash and self.is_mounted:
356
+ try:
357
+ splash = self.query_one("#splash_screen")
358
+ splash.remove()
359
+ except ValueError:
360
+ pass
361
+
362
+ main_container = Vertical(id="main_container")
363
+
364
+ self.mount(main_container)
365
+
366
+ content_container = Horizontal(id="content_container")
367
+ main_container.mount(content_container)
368
+
369
+ chat_area_container = Vertical(id="chat_area_container")
370
+
371
+ chat_display = Static("", id="chat_display")
372
+ chat_history = VerticalScroll(chat_display, id="chat_history")
373
+ chat_history.can_focus = True
374
+
375
+ status_text = Static("", id="status_text")
376
+ keymap_indicator = Static("", id="keymap_indicator")
377
+
378
+ agent_status_display = Horizontal(
379
+ status_text, keymap_indicator, id="agent_status_display", classes="hidden"
380
+ )
381
+
382
+ chat_prompt = Static("> ", id="chat_prompt")
383
+ chat_input = ChatTextArea(
384
+ "",
385
+ id="chat_input",
386
+ show_line_numbers=False,
387
+ )
388
+ chat_input.set_app_reference(self)
389
+ chat_input_container = Horizontal(chat_prompt, chat_input, id="chat_input_container")
390
+
391
+ agents_tree = Tree("🤖 Active Agents", id="agents_tree")
392
+ agents_tree.root.expand()
393
+ agents_tree.show_root = False
394
+
395
+ agents_tree.show_guide = True
396
+ agents_tree.guide_depth = 3
397
+ agents_tree.guide_style = "dashed"
398
+
399
+ stats_display = Static("", id="stats_display")
400
+
401
+ sidebar = Vertical(agents_tree, stats_display, id="sidebar")
402
+
403
+ content_container.mount(chat_area_container)
404
+ content_container.mount(sidebar)
405
+
406
+ chat_area_container.mount(chat_history)
407
+ chat_area_container.mount(agent_status_display)
408
+ chat_area_container.mount(chat_input_container)
409
+
410
+ self.call_after_refresh(self._focus_chat_input)
411
+
412
+ def _focus_chat_input(self) -> None:
413
+ if len(self.screen_stack) > 1 or self.show_splash:
414
+ return
415
+
416
+ if not self.is_mounted:
417
+ return
418
+
419
+ try:
420
+ chat_input = self.query_one("#chat_input", ChatTextArea)
421
+ chat_input.show_vertical_scrollbar = False
422
+ chat_input.show_horizontal_scrollbar = False
423
+ chat_input.focus()
424
+ except (ValueError, Exception):
425
+ self.call_after_refresh(self._focus_chat_input)
426
+
427
+ def _focus_agents_tree(self) -> None:
428
+ if len(self.screen_stack) > 1 or self.show_splash:
429
+ return
430
+
431
+ if not self.is_mounted:
432
+ return
433
+
434
+ try:
435
+ agents_tree = self.query_one("#agents_tree", Tree)
436
+ agents_tree.focus()
437
+
438
+ if agents_tree.root.children:
439
+ first_node = agents_tree.root.children[0]
440
+ agents_tree.select_node(first_node)
441
+ except (ValueError, Exception):
442
+ self.call_after_refresh(self._focus_agents_tree)
443
+
444
+ def on_mount(self) -> None:
445
+ self.title = "aiptx"
446
+
447
+ self.set_timer(4.5, self._hide_splash_screen)
448
+
449
+ def _hide_splash_screen(self) -> None:
450
+ self.show_splash = False
451
+
452
+ self._start_scan_thread()
453
+
454
+ self.set_interval(0.5, self._update_ui_from_tracer)
455
+
456
+ def _update_ui_from_tracer(self) -> None:
457
+ if self.show_splash:
458
+ return
459
+
460
+ if len(self.screen_stack) > 1:
461
+ return
462
+
463
+ if not self.is_mounted:
464
+ return
465
+
466
+ try:
467
+ chat_history = self.query_one("#chat_history", VerticalScroll)
468
+ agents_tree = self.query_one("#agents_tree", Tree)
469
+
470
+ if not self._is_widget_safe(chat_history) or not self._is_widget_safe(agents_tree):
471
+ return
472
+ except (ValueError, Exception):
473
+ return
474
+
475
+ agent_updates = False
476
+ for agent_id, agent_data in self.tracer.agents.items():
477
+ if agent_id not in self._displayed_agents:
478
+ self._add_agent_node(agent_data)
479
+ self._displayed_agents.add(agent_id)
480
+ agent_updates = True
481
+ elif self._update_agent_node(agent_id, agent_data):
482
+ agent_updates = True
483
+
484
+ if agent_updates:
485
+ self._expand_all_agent_nodes()
486
+
487
+ self._update_chat_view()
488
+
489
+ self._update_agent_status_display()
490
+
491
+ self._update_stats_display()
492
+
493
+ def _update_agent_node(self, agent_id: str, agent_data: dict[str, Any]) -> bool:
494
+ if agent_id not in self.agent_nodes:
495
+ return False
496
+
497
+ try:
498
+ agent_node = self.agent_nodes[agent_id]
499
+ agent_name_raw = agent_data.get("name", "Agent")
500
+ status = agent_data.get("status", "running")
501
+
502
+ status_indicators = {
503
+ "running": "🟢",
504
+ "waiting": "⏸️",
505
+ "completed": "✅",
506
+ "failed": "❌",
507
+ "stopped": "⏹️",
508
+ "stopping": "⏸️",
509
+ "llm_failed": "🔴",
510
+ }
511
+
512
+ status_icon = status_indicators.get(status, "🔵")
513
+ agent_name = f"{status_icon} {escape_markup(agent_name_raw)}"
514
+
515
+ if status == "running":
516
+ self._start_agent_verb_timer(agent_id)
517
+ elif status == "waiting":
518
+ self._stop_agent_verb_timer(agent_id)
519
+ else:
520
+ self._stop_agent_verb_timer(agent_id)
521
+
522
+ if agent_node.label != agent_name:
523
+ agent_node.set_label(agent_name)
524
+ return True
525
+
526
+ except (KeyError, AttributeError, ValueError) as e:
527
+ import logging
528
+
529
+ logging.warning(f"Failed to update agent node label: {e}")
530
+
531
+ return False
532
+
533
+ def _update_chat_view(self) -> None:
534
+ if len(self.screen_stack) > 1 or self.show_splash:
535
+ return
536
+
537
+ if not self.is_mounted:
538
+ return
539
+
540
+ try:
541
+ chat_history = self.query_one("#chat_history", VerticalScroll)
542
+ except (ValueError, Exception):
543
+ return
544
+
545
+ if not self._is_widget_safe(chat_history):
546
+ return
547
+
548
+ try:
549
+ is_at_bottom = chat_history.scroll_y >= chat_history.max_scroll_y
550
+ except (AttributeError, ValueError):
551
+ is_at_bottom = True
552
+
553
+ if not self.selected_agent_id:
554
+ content, css_class = self._get_chat_placeholder_content(
555
+ "Select an agent from the tree to see its activity.", "placeholder-no-agent"
556
+ )
557
+ else:
558
+ events = self._gather_agent_events(self.selected_agent_id)
559
+ if not events:
560
+ content, css_class = self._get_chat_placeholder_content(
561
+ "Starting agent...", "placeholder-no-activity"
562
+ )
563
+ else:
564
+ current_event_ids = [e["id"] for e in events]
565
+ if current_event_ids == self._displayed_events:
566
+ return
567
+ content = self._get_rendered_events_content(events)
568
+ css_class = "chat-content"
569
+ self._displayed_events = current_event_ids
570
+
571
+ chat_display = self.query_one("#chat_display", Static)
572
+ self._update_static_content_safe(chat_display, content)
573
+
574
+ chat_display.set_classes(css_class)
575
+
576
+ if is_at_bottom:
577
+ self.call_later(chat_history.scroll_end, animate=False)
578
+
579
+ def _get_chat_placeholder_content(
580
+ self, message: str, placeholder_class: str
581
+ ) -> tuple[str, str]:
582
+ self._displayed_events = [placeholder_class]
583
+ return message, f"chat-placeholder {placeholder_class}"
584
+
585
+ def _get_rendered_events_content(self, events: list[dict[str, Any]]) -> str:
586
+ if not events:
587
+ return ""
588
+
589
+ content_lines = []
590
+ for event in events:
591
+ if event["type"] == "chat":
592
+ chat_content = self._render_chat_content(event["data"])
593
+ if chat_content:
594
+ content_lines.append(chat_content)
595
+ elif event["type"] == "tool":
596
+ tool_content = self._render_tool_content_simple(event["data"])
597
+ if tool_content:
598
+ content_lines.append(tool_content)
599
+
600
+ return "\n\n".join(content_lines)
601
+
602
+ def _update_agent_status_display(self) -> None:
603
+ try:
604
+ status_display = self.query_one("#agent_status_display", Horizontal)
605
+ status_text = self.query_one("#status_text", Static)
606
+ keymap_indicator = self.query_one("#keymap_indicator", Static)
607
+ except (ValueError, Exception):
608
+ return
609
+
610
+ widgets = [status_display, status_text, keymap_indicator]
611
+ if not all(self._is_widget_safe(w) for w in widgets):
612
+ return
613
+
614
+ if not self.selected_agent_id:
615
+ self._safe_widget_operation(status_display.add_class, "hidden")
616
+ return
617
+
618
+ try:
619
+ agent_data = self.tracer.agents[self.selected_agent_id]
620
+ status = agent_data.get("status", "running")
621
+
622
+ if status == "stopping":
623
+ self._safe_widget_operation(status_text.update, "Agent stopping...")
624
+ self._safe_widget_operation(keymap_indicator.update, "")
625
+ self._safe_widget_operation(status_display.remove_class, "hidden")
626
+ elif status == "stopped":
627
+ self._safe_widget_operation(status_text.update, "Agent stopped")
628
+ self._safe_widget_operation(keymap_indicator.update, "")
629
+ self._safe_widget_operation(status_display.remove_class, "hidden")
630
+ elif status == "completed":
631
+ self._safe_widget_operation(status_text.update, "Agent completed")
632
+ self._safe_widget_operation(keymap_indicator.update, "")
633
+ self._safe_widget_operation(status_display.remove_class, "hidden")
634
+ elif status == "llm_failed":
635
+ error_msg = agent_data.get("error_message", "")
636
+ display_msg = (
637
+ f"[red]{escape_markup(error_msg)}[/red]"
638
+ if error_msg
639
+ else "[red]LLM request failed[/red]"
640
+ )
641
+ self._safe_widget_operation(status_text.update, display_msg)
642
+ self._safe_widget_operation(
643
+ keymap_indicator.update, "[dim]Send message to retry[/dim]"
644
+ )
645
+ self._safe_widget_operation(status_display.remove_class, "hidden")
646
+ self._stop_dot_animation()
647
+ elif status == "waiting":
648
+ animated_text = self._get_animated_waiting_text(self.selected_agent_id)
649
+ self._safe_widget_operation(status_text.update, animated_text)
650
+ self._safe_widget_operation(
651
+ keymap_indicator.update, "[dim]Send message to resume[/dim]"
652
+ )
653
+ self._safe_widget_operation(status_display.remove_class, "hidden")
654
+ self._start_dot_animation()
655
+ elif status == "running":
656
+ current_verb = self._get_agent_verb(self.selected_agent_id)
657
+ animated_text = self._get_animated_verb_text(self.selected_agent_id, current_verb)
658
+ self._safe_widget_operation(status_text.update, animated_text)
659
+ self._safe_widget_operation(
660
+ keymap_indicator.update, "[dim]ESC to stop | CTRL-C to quit and save[/dim]"
661
+ )
662
+ self._safe_widget_operation(status_display.remove_class, "hidden")
663
+ self._start_dot_animation()
664
+ else:
665
+ self._safe_widget_operation(status_display.add_class, "hidden")
666
+
667
+ except (KeyError, Exception):
668
+ self._safe_widget_operation(status_display.add_class, "hidden")
669
+
670
+ def _update_stats_display(self) -> None:
671
+ try:
672
+ stats_display = self.query_one("#stats_display", Static)
673
+ except (ValueError, Exception):
674
+ return
675
+
676
+ if not self._is_widget_safe(stats_display):
677
+ return
678
+
679
+ stats_content = Text()
680
+
681
+ stats_text = build_live_stats_text(self.tracer, self.agent_config)
682
+ if stats_text:
683
+ stats_content.append(stats_text)
684
+
685
+ from rich.panel import Panel
686
+
687
+ stats_panel = Panel(
688
+ stats_content,
689
+ title="📊 Live Stats",
690
+ title_align="left",
691
+ border_style="#22c55e",
692
+ padding=(0, 1),
693
+ )
694
+
695
+ self._safe_widget_operation(stats_display.update, stats_panel)
696
+
697
+ def _get_agent_verb(self, agent_id: str) -> str:
698
+ if agent_id not in self._agent_verbs:
699
+ self._agent_verbs[agent_id] = random.choice(self._action_verbs) # nosec B311 # noqa: S311
700
+ return self._agent_verbs[agent_id]
701
+
702
+ def _start_agent_verb_timer(self, agent_id: str) -> None:
703
+ if agent_id not in self._agent_verb_timers:
704
+ self._agent_verb_timers[agent_id] = self.set_interval(
705
+ 30.0, lambda: self._change_agent_action_verb(agent_id)
706
+ )
707
+
708
+ def _stop_agent_verb_timer(self, agent_id: str) -> None:
709
+ if agent_id in self._agent_verb_timers:
710
+ self._agent_verb_timers[agent_id].stop()
711
+ del self._agent_verb_timers[agent_id]
712
+
713
+ def _change_agent_action_verb(self, agent_id: str) -> None:
714
+ if agent_id not in self._agent_verbs:
715
+ self._agent_verbs[agent_id] = random.choice(self._action_verbs) # nosec B311 # noqa: S311
716
+ return
717
+
718
+ current_verb = self._agent_verbs[agent_id]
719
+ available_verbs = [verb for verb in self._action_verbs if verb != current_verb]
720
+ self._agent_verbs[agent_id] = random.choice(available_verbs) # nosec B311 # noqa: S311
721
+
722
+ if self.selected_agent_id == agent_id:
723
+ self._update_agent_status_display()
724
+
725
+ def _get_animated_verb_text(self, agent_id: str, verb: str) -> str:
726
+ if agent_id not in self._agent_dot_states:
727
+ self._agent_dot_states[agent_id] = 0
728
+
729
+ dot_count = self._agent_dot_states[agent_id]
730
+ dots = "." * dot_count
731
+ return f"{verb}{dots}"
732
+
733
+ def _get_animated_waiting_text(self, agent_id: str) -> str:
734
+ if agent_id not in self._agent_dot_states:
735
+ self._agent_dot_states[agent_id] = 0
736
+
737
+ dot_count = self._agent_dot_states[agent_id]
738
+ dots = "." * dot_count
739
+
740
+ return f"Waiting{dots}"
741
+
742
+ def _start_dot_animation(self) -> None:
743
+ if self._dot_animation_timer is None:
744
+ self._dot_animation_timer = self.set_interval(0.6, self._animate_dots)
745
+
746
+ def _stop_dot_animation(self) -> None:
747
+ if self._dot_animation_timer is not None:
748
+ self._dot_animation_timer.stop()
749
+ self._dot_animation_timer = None
750
+
751
+ def _animate_dots(self) -> None:
752
+ has_active_agents = False
753
+
754
+ for agent_id, agent_data in self.tracer.agents.items():
755
+ status = agent_data.get("status", "running")
756
+ if status in ["running", "waiting"]:
757
+ has_active_agents = True
758
+ current_dots = self._agent_dot_states.get(agent_id, 0)
759
+ self._agent_dot_states[agent_id] = (current_dots + 1) % 4
760
+
761
+ if (
762
+ has_active_agents
763
+ and self.selected_agent_id
764
+ and self.selected_agent_id in self.tracer.agents
765
+ ):
766
+ selected_status = self.tracer.agents[self.selected_agent_id].get("status", "running")
767
+ if selected_status in ["running", "waiting"]:
768
+ self._update_agent_status_display()
769
+
770
+ if not has_active_agents:
771
+ self._stop_dot_animation()
772
+ for agent_id in list(self._agent_dot_states.keys()):
773
+ if agent_id not in self.tracer.agents or self.tracer.agents[agent_id].get(
774
+ "status"
775
+ ) not in ["running", "waiting"]:
776
+ del self._agent_dot_states[agent_id]
777
+
778
+ def _gather_agent_events(self, agent_id: str) -> list[dict[str, Any]]:
779
+ chat_events = [
780
+ {
781
+ "type": "chat",
782
+ "timestamp": msg["timestamp"],
783
+ "id": f"chat_{msg['message_id']}",
784
+ "data": msg,
785
+ }
786
+ for msg in self.tracer.chat_messages
787
+ if msg.get("agent_id") == agent_id
788
+ ]
789
+
790
+ tool_events = [
791
+ {
792
+ "type": "tool",
793
+ "timestamp": tool_data["timestamp"],
794
+ "id": f"tool_{exec_id}",
795
+ "data": tool_data,
796
+ }
797
+ for exec_id, tool_data in self.tracer.tool_executions.items()
798
+ if tool_data.get("agent_id") == agent_id
799
+ ]
800
+
801
+ events = chat_events + tool_events
802
+ events.sort(key=lambda e: (e["timestamp"], e["id"]))
803
+ return events
804
+
805
+ def watch_selected_agent_id(self, _agent_id: str | None) -> None:
806
+ if len(self.screen_stack) > 1 or self.show_splash:
807
+ return
808
+
809
+ if not self.is_mounted:
810
+ return
811
+
812
+ self._displayed_events.clear()
813
+
814
+ self.call_later(self._update_chat_view)
815
+ self._update_agent_status_display()
816
+
817
+ def _start_scan_thread(self) -> None:
818
+ def scan_target() -> None:
819
+ try:
820
+ loop = asyncio.new_event_loop()
821
+ asyncio.set_event_loop(loop)
822
+
823
+ try:
824
+ agent = AIPTxAgent(self.agent_config)
825
+
826
+ if not self._scan_stop_event.is_set():
827
+ loop.run_until_complete(agent.execute_scan(self.scan_config))
828
+
829
+ except (KeyboardInterrupt, asyncio.CancelledError):
830
+ logging.info("Scan interrupted by user")
831
+ except (ConnectionError, TimeoutError):
832
+ logging.exception("Network error during scan")
833
+ except RuntimeError:
834
+ logging.exception("Runtime error during scan")
835
+ except Exception:
836
+ logging.exception("Unexpected error during scan")
837
+ finally:
838
+ loop.close()
839
+ self._scan_completed.set()
840
+
841
+ except Exception:
842
+ logging.exception("Error setting up scan thread")
843
+ self._scan_completed.set()
844
+
845
+ self._scan_thread = threading.Thread(target=scan_target, daemon=True)
846
+ self._scan_thread.start()
847
+
848
+ def _add_agent_node(self, agent_data: dict[str, Any]) -> None:
849
+ if len(self.screen_stack) > 1 or self.show_splash:
850
+ return
851
+
852
+ if not self.is_mounted:
853
+ return
854
+
855
+ agent_id = agent_data["id"]
856
+ parent_id = agent_data.get("parent_id")
857
+ status = agent_data.get("status", "running")
858
+
859
+ try:
860
+ agents_tree = self.query_one("#agents_tree", Tree)
861
+ except (ValueError, Exception):
862
+ return
863
+
864
+ agent_name_raw = agent_data.get("name", "Agent")
865
+
866
+ status_indicators = {
867
+ "running": "🟢",
868
+ "waiting": "🟡",
869
+ "completed": "✅",
870
+ "failed": "❌",
871
+ "stopped": "⏹️",
872
+ "stopping": "⏸️",
873
+ }
874
+
875
+ status_icon = status_indicators.get(status, "🔵")
876
+ agent_name = f"{status_icon} {escape_markup(agent_name_raw)}"
877
+
878
+ if status in ["running", "waiting"]:
879
+ self._start_agent_verb_timer(agent_id)
880
+
881
+ try:
882
+ if parent_id and parent_id in self.agent_nodes:
883
+ parent_node = self.agent_nodes[parent_id]
884
+ agent_node = parent_node.add(
885
+ agent_name,
886
+ data={"agent_id": agent_id},
887
+ )
888
+ parent_node.allow_expand = True
889
+ else:
890
+ agent_node = agents_tree.root.add(
891
+ agent_name,
892
+ data={"agent_id": agent_id},
893
+ )
894
+
895
+ agent_node.allow_expand = False
896
+ agent_node.expand()
897
+ self.agent_nodes[agent_id] = agent_node
898
+
899
+ if len(self.agent_nodes) == 1:
900
+ agents_tree.select_node(agent_node)
901
+ self.selected_agent_id = agent_id
902
+
903
+ self._reorganize_orphaned_agents(agent_id)
904
+ except (AttributeError, ValueError, RuntimeError) as e:
905
+ import logging
906
+
907
+ logging.warning(f"Failed to add agent node {agent_id}: {e}")
908
+
909
+ def _expand_all_agent_nodes(self) -> None:
910
+ if len(self.screen_stack) > 1 or self.show_splash:
911
+ return
912
+
913
+ if not self.is_mounted:
914
+ return
915
+
916
+ try:
917
+ agents_tree = self.query_one("#agents_tree", Tree)
918
+ self._expand_node_recursively(agents_tree.root)
919
+ except (ValueError, Exception):
920
+ logging.debug("Tree not ready for expanding nodes")
921
+
922
+ def _expand_node_recursively(self, node: TreeNode) -> None:
923
+ if not node.is_expanded:
924
+ node.expand()
925
+ for child in node.children:
926
+ self._expand_node_recursively(child)
927
+
928
+ def _copy_node_under(self, node_to_copy: TreeNode, new_parent: TreeNode) -> None:
929
+ agent_id = node_to_copy.data["agent_id"]
930
+ agent_data = self.tracer.agents.get(agent_id, {})
931
+ agent_name_raw = agent_data.get("name", "Agent")
932
+ status = agent_data.get("status", "running")
933
+
934
+ status_indicators = {
935
+ "running": "🟢",
936
+ "waiting": "🟡",
937
+ "completed": "✅",
938
+ "failed": "❌",
939
+ "stopped": "⏹️",
940
+ "stopping": "⏸️",
941
+ }
942
+
943
+ status_icon = status_indicators.get(status, "🔵")
944
+ agent_name = f"{status_icon} {escape_markup(agent_name_raw)}"
945
+
946
+ new_node = new_parent.add(
947
+ agent_name,
948
+ data=node_to_copy.data,
949
+ )
950
+ new_node.allow_expand = node_to_copy.allow_expand
951
+
952
+ self.agent_nodes[agent_id] = new_node
953
+
954
+ for child in node_to_copy.children:
955
+ self._copy_node_under(child, new_node)
956
+
957
+ if node_to_copy.is_expanded:
958
+ new_node.expand()
959
+
960
+ def _reorganize_orphaned_agents(self, new_parent_id: str) -> None:
961
+ agents_to_move = []
962
+
963
+ for agent_id, agent_data in self.tracer.agents.items():
964
+ if (
965
+ agent_data.get("parent_id") == new_parent_id
966
+ and agent_id in self.agent_nodes
967
+ and agent_id != new_parent_id
968
+ ):
969
+ agents_to_move.append(agent_id)
970
+
971
+ if not agents_to_move:
972
+ return
973
+
974
+ parent_node = self.agent_nodes[new_parent_id]
975
+
976
+ for child_agent_id in agents_to_move:
977
+ if child_agent_id in self.agent_nodes:
978
+ old_node = self.agent_nodes[child_agent_id]
979
+
980
+ if old_node.parent is parent_node:
981
+ continue
982
+
983
+ self._copy_node_under(old_node, parent_node)
984
+
985
+ old_node.remove()
986
+
987
+ parent_node.allow_expand = True
988
+ self._expand_all_agent_nodes()
989
+
990
+ def _render_chat_content(self, msg_data: dict[str, Any]) -> str:
991
+ role = msg_data.get("role")
992
+ content = escape_markup(msg_data.get("content", ""))
993
+
994
+ if not content:
995
+ return ""
996
+
997
+ if role == "user":
998
+ from interface.tool_components.user_message_renderer import UserMessageRenderer
999
+
1000
+ return UserMessageRenderer.render_simple(content)
1001
+ return content
1002
+
1003
+ def _render_tool_content_simple(self, tool_data: dict[str, Any]) -> str:
1004
+ tool_name = tool_data.get("tool_name", "Unknown Tool")
1005
+ args = tool_data.get("args", {})
1006
+ status = tool_data.get("status", "unknown")
1007
+ result = tool_data.get("result")
1008
+
1009
+ tool_colors = {
1010
+ "terminal_execute": "#22c55e",
1011
+ "browser_action": "#06b6d4",
1012
+ "python_action": "#3b82f6",
1013
+ "agents_graph_action": "#fbbf24",
1014
+ "file_edit_action": "#10b981",
1015
+ "proxy_action": "#06b6d4",
1016
+ "notes_action": "#fbbf24",
1017
+ "thinking_action": "#a855f7",
1018
+ "web_search_action": "#22c55e",
1019
+ "finish_action": "#dc2626",
1020
+ "reporting_action": "#ea580c",
1021
+ "scan_start_info": "#22c55e",
1022
+ "subagent_start_info": "#22c55e",
1023
+ "llm_error_details": "#dc2626",
1024
+ }
1025
+
1026
+ color = tool_colors.get(tool_name, "#737373")
1027
+
1028
+ from interface.tool_components.registry import get_tool_renderer
1029
+
1030
+ renderer = get_tool_renderer(tool_name)
1031
+
1032
+ if renderer:
1033
+ widget = renderer.render(tool_data)
1034
+ content = str(widget.renderable)
1035
+ elif tool_name == "llm_error_details":
1036
+ lines = ["[red]✗ LLM Request Failed[/red]"]
1037
+ if args.get("details"):
1038
+ details = args["details"]
1039
+ if len(details) > 300:
1040
+ details = details[:297] + "..."
1041
+ lines.append(f"[dim]Details:[/dim] {escape_markup(details)}")
1042
+ content = "\n".join(lines)
1043
+ else:
1044
+ status_icons = {
1045
+ "running": "[yellow]●[/yellow]",
1046
+ "completed": "[green]✓[/green]",
1047
+ "failed": "[red]✗[/red]",
1048
+ "error": "[red]✗[/red]",
1049
+ }
1050
+ status_icon = status_icons.get(status, "[dim]○[/dim]")
1051
+
1052
+ lines = [f"→ Using tool [bold blue]{escape_markup(tool_name)}[/] {status_icon}"]
1053
+
1054
+ if args:
1055
+ for k, v in list(args.items())[:2]:
1056
+ str_v = str(v)
1057
+ if len(str_v) > 80:
1058
+ str_v = str_v[:77] + "..."
1059
+ lines.append(f" [dim]{k}:[/] {escape_markup(str_v)}")
1060
+
1061
+ if status in ["completed", "failed", "error"] and result:
1062
+ result_str = str(result)
1063
+ if len(result_str) > 150:
1064
+ result_str = result_str[:147] + "..."
1065
+ lines.append(f"[bold]Result:[/] {escape_markup(result_str)}")
1066
+
1067
+ content = "\n".join(lines)
1068
+
1069
+ lines = content.split("\n")
1070
+ bordered_lines = [f"[{color}]▍[/{color}] {line}" for line in lines]
1071
+ return "\n".join(bordered_lines)
1072
+
1073
+ @on(Tree.NodeHighlighted) # type: ignore[misc]
1074
+ def handle_tree_highlight(self, event: Tree.NodeHighlighted) -> None:
1075
+ if len(self.screen_stack) > 1 or self.show_splash:
1076
+ return
1077
+
1078
+ if not self.is_mounted:
1079
+ return
1080
+
1081
+ node = event.node
1082
+
1083
+ try:
1084
+ agents_tree = self.query_one("#agents_tree", Tree)
1085
+ except (ValueError, Exception):
1086
+ return
1087
+
1088
+ if self.focused == agents_tree and node.data:
1089
+ agent_id = node.data.get("agent_id")
1090
+ if agent_id:
1091
+ self.selected_agent_id = agent_id
1092
+
1093
+ def _send_user_message(self, message: str) -> None:
1094
+ if not self.selected_agent_id:
1095
+ return
1096
+
1097
+ if self.tracer:
1098
+ self.tracer.log_chat_message(
1099
+ content=message,
1100
+ role="user",
1101
+ agent_id=self.selected_agent_id,
1102
+ )
1103
+
1104
+ try:
1105
+ from aipt_v2.tools.agents_graph.agents_graph_actions import send_user_message_to_agent
1106
+
1107
+ send_user_message_to_agent(self.selected_agent_id, message)
1108
+
1109
+ except (ImportError, AttributeError) as e:
1110
+ import logging
1111
+
1112
+ logging.warning(f"Failed to send message to agent {self.selected_agent_id}: {e}")
1113
+
1114
+ self._displayed_events.clear()
1115
+ self._update_chat_view()
1116
+
1117
+ self.call_after_refresh(self._focus_chat_input)
1118
+
1119
+ def _get_agent_name(self, agent_id: str) -> str:
1120
+ try:
1121
+ if self.tracer and agent_id in self.tracer.agents:
1122
+ agent_name = self.tracer.agents[agent_id].get("name")
1123
+ if isinstance(agent_name, str):
1124
+ return agent_name
1125
+ except (KeyError, AttributeError) as e:
1126
+ logging.warning(f"Could not retrieve agent name for {agent_id}: {e}")
1127
+ return "Unknown Agent"
1128
+
1129
+ def action_toggle_help(self) -> None:
1130
+ if self.show_splash or not self.is_mounted:
1131
+ return
1132
+
1133
+ try:
1134
+ self.query_one("#main_container")
1135
+ except (ValueError, Exception):
1136
+ return
1137
+
1138
+ if isinstance(self.screen, HelpScreen):
1139
+ self.pop_screen()
1140
+ return
1141
+
1142
+ if len(self.screen_stack) > 1:
1143
+ return
1144
+
1145
+ self.push_screen(HelpScreen())
1146
+
1147
+ def action_request_quit(self) -> None:
1148
+ if self.show_splash or not self.is_mounted:
1149
+ self.action_custom_quit()
1150
+ return
1151
+
1152
+ if len(self.screen_stack) > 1:
1153
+ return
1154
+
1155
+ try:
1156
+ self.query_one("#main_container")
1157
+ except (ValueError, Exception):
1158
+ self.action_custom_quit()
1159
+ return
1160
+
1161
+ self.push_screen(QuitScreen())
1162
+
1163
+ def action_stop_selected_agent(self) -> None:
1164
+ if (
1165
+ self.show_splash
1166
+ or not self.is_mounted
1167
+ or len(self.screen_stack) > 1
1168
+ or not self.selected_agent_id
1169
+ ):
1170
+ return
1171
+
1172
+ agent_name, should_stop = self._validate_agent_for_stopping()
1173
+ if not should_stop:
1174
+ return
1175
+
1176
+ try:
1177
+ self.query_one("#main_container")
1178
+ except (ValueError, Exception):
1179
+ return
1180
+
1181
+ self.push_screen(StopAgentScreen(agent_name, self.selected_agent_id))
1182
+
1183
+ def _validate_agent_for_stopping(self) -> tuple[str, bool]:
1184
+ agent_name = "Unknown Agent"
1185
+
1186
+ try:
1187
+ if self.tracer and self.selected_agent_id in self.tracer.agents:
1188
+ agent_data = self.tracer.agents[self.selected_agent_id]
1189
+ agent_name = agent_data.get("name", "Unknown Agent")
1190
+
1191
+ agent_status = agent_data.get("status", "running")
1192
+ if agent_status not in ["running"]:
1193
+ return agent_name, False
1194
+
1195
+ agent_events = self._gather_agent_events(self.selected_agent_id)
1196
+ if not agent_events:
1197
+ return agent_name, False
1198
+
1199
+ return agent_name, True
1200
+
1201
+ except (KeyError, AttributeError, ValueError) as e:
1202
+ import logging
1203
+
1204
+ logging.warning(f"Failed to gather agent events: {e}")
1205
+
1206
+ return agent_name, False
1207
+
1208
+ def action_confirm_stop_agent(self, agent_id: str) -> None:
1209
+ self.pop_screen()
1210
+
1211
+ try:
1212
+ from aipt_v2.tools.agents_graph.agents_graph_actions import stop_agent
1213
+
1214
+ result = stop_agent(agent_id)
1215
+
1216
+ import logging
1217
+
1218
+ if result.get("success"):
1219
+ logging.info(f"Stop request sent to agent: {result.get('message', 'Unknown')}")
1220
+ else:
1221
+ logging.warning(f"Failed to stop agent: {result.get('error', 'Unknown error')}")
1222
+
1223
+ except Exception:
1224
+ import logging
1225
+
1226
+ logging.exception(f"Failed to stop agent {agent_id}")
1227
+
1228
+ def action_custom_quit(self) -> None:
1229
+ for agent_id in list(self._agent_verb_timers.keys()):
1230
+ self._stop_agent_verb_timer(agent_id)
1231
+
1232
+ if self._scan_thread and self._scan_thread.is_alive():
1233
+ self._scan_stop_event.set()
1234
+
1235
+ self._scan_thread.join(timeout=1.0)
1236
+
1237
+ self.tracer.cleanup()
1238
+
1239
+ self.exit()
1240
+
1241
+ def _is_widget_safe(self, widget: Any) -> bool:
1242
+ try:
1243
+ _ = widget.screen
1244
+ except (AttributeError, ValueError, Exception):
1245
+ return False
1246
+ else:
1247
+ return bool(widget.is_mounted)
1248
+
1249
+ def _safe_widget_operation(
1250
+ self, operation: Callable[..., Any], *args: Any, **kwargs: Any
1251
+ ) -> bool:
1252
+ try:
1253
+ operation(*args, **kwargs)
1254
+ except (AttributeError, ValueError, Exception):
1255
+ return False
1256
+ else:
1257
+ return True
1258
+
1259
+ def _update_static_content_safe(self, widget: Static, content: str) -> None:
1260
+ try:
1261
+ widget.update(content)
1262
+ except Exception: # noqa: BLE001
1263
+ try:
1264
+ safe_text = Text.from_markup(content)
1265
+ widget.update(safe_text)
1266
+ except Exception: # noqa: BLE001
1267
+ import re
1268
+
1269
+ plain_text = re.sub(r"\[.*?\]", "", content)
1270
+ widget.update(plain_text)
1271
+
1272
+
1273
+ async def run_tui(args: argparse.Namespace) -> None:
1274
+ """Run AIPT in interactive TUI mode with textual."""
1275
+ app = AIPTxTUIApp(args)
1276
+ await app.run_async()