mirror of
https://github.com/usestrix/strix.git
synced 2026-08-28 05:25:00 +00:00
Isolate MCP connections per task and surface connection status in the UIs (#1181)
This commit is contained in:
parent
cbb0f57058
commit
717ffc8f4c
37 changed files with 2203 additions and 183 deletions
|
|
@ -255,6 +255,11 @@ ignore = [
|
|||
"strix/tools/notes/tools.py" = ["PLC0415", "TC002"]
|
||||
"strix/tools/finish/tool.py" = ["PLC0415", "TC002"]
|
||||
"strix/tools/reporting/tool.py" = ["PLC0415", "TC002"]
|
||||
# Lazy imports of strix.tools.mcp.client avoid a circular import (client imports
|
||||
# the session module at module load).
|
||||
"strix/tools/mcp/session.py" = ["PLC0415"]
|
||||
# call_mcp is a chain of guard clauses that each return an error string.
|
||||
"strix/tools/mcp/agent_tools.py" = ["PLR0911"]
|
||||
"strix/tools/**/*.py" = [
|
||||
"ARG001", # Unused function argument (tools may have unused args for interface consistency)
|
||||
]
|
||||
|
|
|
|||
|
|
@ -53,18 +53,43 @@ from strix.tools.output_store import (
|
|||
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from agents.mcp import MCPServer
|
||||
from agents.memory import SQLiteSession
|
||||
from agents.result import RunResultBase
|
||||
|
||||
from strix.runtime.status import StatusSink
|
||||
from strix.tools.mcp import ConnectedMcpServer, McpConnectionRequest
|
||||
from strix.tools.mcp import (
|
||||
ConnectedMcpServer,
|
||||
McpConnectionRequest,
|
||||
McpRegistry,
|
||||
SupervisedMcpSession,
|
||||
)
|
||||
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
StreamEventSink = Callable[[str, Any], None]
|
||||
|
||||
# Receives the run's MCP connection roster as a list of non-secret status dicts
|
||||
# ({"name", "provider", "tool_count", "dead"}), once when the connections are
|
||||
# established and again each time a connection transitions to dead. An interface
|
||||
# can persist it, render it, or forward it on as connection status. Kept as a
|
||||
# snapshot of the whole roster (not a per-
|
||||
# connection delta) so every call carries a consistent, current picture.
|
||||
McpStatusSink = Callable[[list[dict[str, Any]]], None]
|
||||
|
||||
|
||||
def _mcp_roster_payload(registry: McpRegistry) -> list[dict[str, Any]]:
|
||||
"""The run's MCP roster as non-secret status dicts (name/provider/tool_count/dead)."""
|
||||
return [
|
||||
{
|
||||
"name": status.name,
|
||||
"provider": status.provider,
|
||||
"tool_count": status.tool_count,
|
||||
"dead": status.dead,
|
||||
}
|
||||
for status in registry.statuses()
|
||||
]
|
||||
|
||||
|
||||
def _mcp_startup_summary(connections: list[ConnectedMcpServer]) -> str:
|
||||
"""One user-facing line summarizing the MCP servers that connected."""
|
||||
|
|
@ -91,6 +116,21 @@ def _record_mcp_connections(connections: list[ConnectedMcpServer]) -> None:
|
|||
report_state.record_mcp_connections([connection.name for connection in connections])
|
||||
|
||||
|
||||
def _persist_mcp_status(roster: list[dict[str, Any]]) -> None:
|
||||
"""Write the run's non-secret MCP connection status roster to run.json.
|
||||
|
||||
The viewer rebuilds its display by re-reading the run's files from disk, so
|
||||
it cannot see the in-memory ``mcp_status_sink`` the TUI consumes. Persisting
|
||||
the same non-secret roster (name / provider / tool_count / dead) gives the
|
||||
viewer a source it can poll. Runs regardless of whether an interface sink is
|
||||
attached, so the standalone / non-TUI CLI path records health too.
|
||||
"""
|
||||
report_state = get_global_report_state()
|
||||
if report_state is None:
|
||||
return
|
||||
report_state.record_mcp_connection_status(roster)
|
||||
|
||||
|
||||
def _merge_root_prompt_context(
|
||||
scope_context: dict[str, Any],
|
||||
extra_system_prompt_context: dict[str, Any] | None,
|
||||
|
|
@ -157,6 +197,7 @@ async def run_strix_scan(
|
|||
extra_system_prompt_context: dict[str, Any] | None = None,
|
||||
status_sink: StatusSink | None = None,
|
||||
mcp_connection_requests: list[McpConnectionRequest] | None = None,
|
||||
mcp_status_sink: McpStatusSink | None = None,
|
||||
) -> RunResultBase | None:
|
||||
"""Run or resume one Strix scan against a sandbox.
|
||||
|
||||
|
|
@ -297,7 +338,7 @@ async def run_strix_scan(
|
|||
configure_spill_writer(_spill_to_workspace)
|
||||
|
||||
sessions_to_close: list[SQLiteSession] = []
|
||||
mcp_servers: list[MCPServer] = []
|
||||
mcp_sessions: list[SupervisedMcpSession] = []
|
||||
|
||||
try:
|
||||
targets = scan_config.get("targets") or []
|
||||
|
|
@ -366,7 +407,7 @@ async def run_strix_scan(
|
|||
mcp_requests = mcp_connection_requests
|
||||
if mcp_requests:
|
||||
connections = await attach_mcp_requests(mcp_requests, mcp_registry)
|
||||
mcp_servers = [c.server for c in connections]
|
||||
mcp_sessions = [c.session for c in connections]
|
||||
# Recorded even when nothing connected, so a resumed run does not
|
||||
# keep attributing tool calls to servers it no longer has.
|
||||
_record_mcp_connections(connections)
|
||||
|
|
@ -387,6 +428,31 @@ async def run_strix_scan(
|
|||
}
|
||||
for summary in mcp_registry.summaries()
|
||||
]
|
||||
# Feed a non-secret connection roster (name / provider /
|
||||
# tool_count / dead) to two consumers: once now (all
|
||||
# currently healthy) and again whenever a connection later
|
||||
# dies. It is always persisted to run.json so the viewer,
|
||||
# which re-reads the run's files from disk, can render the
|
||||
# MCP connections panel and health without an in-memory
|
||||
# sink. When an interface sink is attached (the TUI backend,
|
||||
# or pro forwarding into the app's event stream) it also
|
||||
# receives the same snapshot. In-use is derived separately by
|
||||
# each interface from the connection-tagged tool-call events,
|
||||
# so it is not carried here.
|
||||
def _emit_mcp_status() -> None:
|
||||
roster = _mcp_roster_payload(mcp_registry)
|
||||
_persist_mcp_status(roster)
|
||||
if mcp_status_sink is not None:
|
||||
try:
|
||||
mcp_status_sink(roster)
|
||||
except Exception:
|
||||
logger.exception("MCP status sink failed")
|
||||
|
||||
for connection_name in mcp_registry.names():
|
||||
entry = mcp_registry.get(connection_name)
|
||||
if entry is not None:
|
||||
entry.session.set_on_dead(_emit_mcp_status)
|
||||
_emit_mcp_status()
|
||||
except Exception:
|
||||
logger.exception("Failed to connect user MCP servers; continuing without them")
|
||||
|
||||
|
|
@ -581,9 +647,9 @@ async def run_strix_scan(
|
|||
for s in sessions_to_close:
|
||||
with contextlib.suppress(Exception):
|
||||
s.close()
|
||||
for mcp_server in mcp_servers:
|
||||
for mcp_session in mcp_sessions:
|
||||
with contextlib.suppress(Exception):
|
||||
await mcp_server.cleanup() # type: ignore[no-untyped-call]
|
||||
await mcp_session.aclose()
|
||||
with contextlib.suppress(Exception):
|
||||
await coordinator._maybe_snapshot()
|
||||
if cleanup_on_exit:
|
||||
|
|
|
|||
|
|
@ -103,6 +103,11 @@ class TuiController:
|
|||
self.messages: list[dict[str, str]] = []
|
||||
self._next_message_id = 1
|
||||
self.error: str | None = None
|
||||
# The run's MCP connection roster (name / tool_count / dead), pushed by
|
||||
# the engine via the mcp_status_sink once the connections are established
|
||||
# and again each time one dies. Empty for a run with no MCP connections,
|
||||
# so the Go sidebar simply omits the panel. Non-secret by construction.
|
||||
self.mcp_connections: list[dict[str, Any]] = []
|
||||
self.viewer_status = "idle"
|
||||
self.viewer_url: str | None = None
|
||||
self._viewer_httpd: Any = None
|
||||
|
|
@ -128,6 +133,24 @@ class TuiController:
|
|||
if scan_loop is not None:
|
||||
self.scan_loop = scan_loop
|
||||
|
||||
def set_mcp_connections(self, roster: list[dict[str, Any]]) -> None:
|
||||
"""Store the run's MCP connection roster and repaint.
|
||||
|
||||
``roster`` is the engine's non-secret status snapshot: one entry per
|
||||
connection carrying ``name``, ``tool_count``, and ``dead``. Called once
|
||||
when the connections are established (all healthy) and again whenever a
|
||||
connection dies (the same whole-roster snapshot, with that one now dead)."""
|
||||
self.mcp_connections = [
|
||||
{
|
||||
"name": str(entry.get("name", "")),
|
||||
"tool_count": int(entry.get("tool_count", 0) or 0),
|
||||
"dead": bool(entry.get("dead", False)),
|
||||
}
|
||||
for entry in roster
|
||||
if isinstance(entry, dict) and entry.get("name")
|
||||
]
|
||||
self.notify_changed()
|
||||
|
||||
def begin_preparation(self) -> None:
|
||||
"""Mark a directly-launched run as preparing behind the live TUI."""
|
||||
self.scan_state = "preparing"
|
||||
|
|
@ -200,6 +223,14 @@ class TuiController:
|
|||
],
|
||||
"usage": terminal_projection(usage, max_string=256, max_items=20),
|
||||
"subscription": subscription,
|
||||
"connections": [
|
||||
{
|
||||
"name": terminal_projection(entry["name"], max_string=64),
|
||||
"tool_count": entry["tool_count"],
|
||||
"dead": entry["dead"],
|
||||
}
|
||||
for entry in self.mcp_connections[:32]
|
||||
],
|
||||
"viewer_status": self.viewer_status,
|
||||
"viewer_url": terminal_projection(self.viewer_url, max_string=1024),
|
||||
"error": terminal_projection(self.error, max_string=2 * 1024),
|
||||
|
|
|
|||
|
|
@ -177,6 +177,7 @@ def bounded_state_projection(state: dict[str, Any]) -> dict[str, Any]:
|
|||
"messages": [],
|
||||
"usage": state["usage"],
|
||||
"subscription": state["subscription"],
|
||||
"connections": state.get("connections", [])[:32],
|
||||
"viewer_status": state["viewer_status"],
|
||||
"viewer_url": None,
|
||||
"error": terminal_projection(state["error"], max_string=256),
|
||||
|
|
|
|||
|
|
@ -209,7 +209,7 @@ func (m *Model) ensureAgentVisible() {
|
|||
m.agentOffset = 0
|
||||
return
|
||||
}
|
||||
_, _, agentHeight := m.sidebarHeights()
|
||||
_, _, _, agentHeight := m.sidebarHeights()
|
||||
rows := max(1, agentHeight-4)
|
||||
row := selectedAgentRow(entries, m.selectedAgent)
|
||||
if row < m.agentOffset {
|
||||
|
|
@ -221,7 +221,7 @@ func (m *Model) ensureAgentVisible() {
|
|||
}
|
||||
|
||||
func (m Model) agentPageSize() int {
|
||||
_, _, agentHeight := m.sidebarHeights()
|
||||
_, _, _, agentHeight := m.sidebarHeights()
|
||||
return max(1, agentHeight-4)
|
||||
}
|
||||
|
||||
|
|
|
|||
105
strix/interface/tui/internal/app/mcp_test.go
Normal file
105
strix/interface/tui/internal/app/mcp_test.go
Normal file
|
|
@ -0,0 +1,105 @@
|
|||
package app
|
||||
|
||||
import (
|
||||
"fmt"
|
||||
"strings"
|
||||
"testing"
|
||||
|
||||
"github.com/charmbracelet/x/ansi"
|
||||
"github.com/usestrix/strix/tui/internal/protocol"
|
||||
)
|
||||
|
||||
func mcpModel(t *testing.T) Model {
|
||||
t.Helper()
|
||||
m := New(nil)
|
||||
m.width, m.height = 130, 40
|
||||
m.showSplash = false
|
||||
m.handleEnvelope(stateEnvelope(t, 1, protocol.Snapshot{
|
||||
ScanState: "running",
|
||||
Connections: []protocol.Connection{
|
||||
{Name: "supabase", ToolCount: 3, Dead: false},
|
||||
{Name: "vercel", ToolCount: 1, Dead: true},
|
||||
},
|
||||
}))
|
||||
return m
|
||||
}
|
||||
|
||||
func TestMcpPanelShowsHealthyAndOffline(t *testing.T) {
|
||||
m := mcpModel(t)
|
||||
out := ansi.Strip(m.mcpConnectionsView(40, 6))
|
||||
for _, want := range []string{"MCP Connections (2)", "supabase", "3 tools", "vercel", "offline"} {
|
||||
if !strings.Contains(out, want) {
|
||||
t.Fatalf("panel missing %q:\n%s", want, out)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// A roster longer than the panel height shows a window of rows rather than every
|
||||
// connection, while the header keeps the full count.
|
||||
func TestMcpPanelWindowsLargeRosterAndCountsAll(t *testing.T) {
|
||||
m := New(nil)
|
||||
m.width, m.height = 130, 40
|
||||
m.showSplash = false
|
||||
conns := make([]protocol.Connection, 0, 12)
|
||||
for i := 0; i < 12; i++ {
|
||||
conns = append(conns, protocol.Connection{Name: fmt.Sprintf("conn-%02d", i), ToolCount: 2})
|
||||
}
|
||||
m.snapshot.Connections = conns
|
||||
|
||||
// rows = 6 → one header line + five roster rows.
|
||||
out := ansi.Strip(m.mcpConnectionsView(40, 6))
|
||||
if !strings.Contains(out, "MCP Connections (12)") {
|
||||
t.Fatalf("header did not carry the full connection count:\n%s", out)
|
||||
}
|
||||
if !strings.Contains(out, "conn-00") {
|
||||
t.Fatalf("top of the roster was not rendered:\n%s", out)
|
||||
}
|
||||
if strings.Contains(out, "conn-11") {
|
||||
t.Fatalf("a roster past the panel height should be windowed, not fully drawn:\n%s", out)
|
||||
}
|
||||
if got := strings.Count(out, "\n") + 1; got != 6 {
|
||||
t.Fatalf("panel rendered %d lines, want 6 (header + five rows)", got)
|
||||
}
|
||||
|
||||
// Scrolling the roster brings the tail into view while the header count holds.
|
||||
m.mcpOffset = 7
|
||||
scrolled := ansi.Strip(m.mcpConnectionsView(40, 6))
|
||||
if !strings.Contains(scrolled, "conn-11") || !strings.Contains(scrolled, "MCP Connections (12)") {
|
||||
t.Fatalf("scrolled window did not reveal the tail with the count intact:\n%s", scrolled)
|
||||
}
|
||||
}
|
||||
|
||||
func TestMcpPanelHeightReservedFromAgentBudget(t *testing.T) {
|
||||
m := mcpModel(t)
|
||||
_, _, mcpHeight, _ := m.sidebarHeights()
|
||||
if mcpHeight <= 0 {
|
||||
t.Fatalf("connections present but no panel height was reserved: %d", mcpHeight)
|
||||
}
|
||||
|
||||
empty := New(nil)
|
||||
empty.width, empty.height = 130, 40
|
||||
empty.showSplash = false
|
||||
empty.handleEnvelope(stateEnvelope(t, 1, protocol.Snapshot{ScanState: "running"}))
|
||||
if _, _, emptyHeight, _ := empty.sidebarHeights(); emptyHeight != 0 {
|
||||
t.Fatalf("no connections should leave the panel absent, got height %d", emptyHeight)
|
||||
}
|
||||
}
|
||||
|
||||
func TestMcpInUseReadsRunningConnectionTaggedCalls(t *testing.T) {
|
||||
m := mcpModel(t)
|
||||
m.handleEnvelope(bootstrapEnvelope(t, "events", 1,
|
||||
protocol.Event{ID: "e1", Type: "tool", AgentID: "a1", Data: map[string]any{
|
||||
"tool_name": "call_mcp", "mcp_connection": "supabase", "status": "running",
|
||||
}},
|
||||
protocol.Event{ID: "e2", Type: "tool", AgentID: "a1", Data: map[string]any{
|
||||
"tool_name": "call_mcp", "mcp_connection": "vercel", "status": "completed",
|
||||
}},
|
||||
))
|
||||
inUse := m.mcpInUse()
|
||||
if !inUse["supabase"] {
|
||||
t.Fatalf("a running connection-tagged call should mark the connection in use")
|
||||
}
|
||||
if inUse["vercel"] {
|
||||
t.Fatalf("a completed call must not mark the connection in use")
|
||||
}
|
||||
}
|
||||
|
|
@ -73,6 +73,7 @@ const (
|
|||
focusChat
|
||||
focusAgents
|
||||
focusVulnerabilities
|
||||
focusMcp
|
||||
)
|
||||
|
||||
type scrollbarTarget int
|
||||
|
|
@ -82,6 +83,7 @@ const (
|
|||
scrollbarTrace
|
||||
scrollbarAgents
|
||||
scrollbarFindings
|
||||
scrollbarMcp
|
||||
)
|
||||
|
||||
type Model struct {
|
||||
|
|
@ -109,6 +111,7 @@ type Model struct {
|
|||
selectedVuln int
|
||||
agentOffset int
|
||||
vulnOffset int
|
||||
mcpOffset int
|
||||
modalChoice int
|
||||
reportFocus string
|
||||
ready bool
|
||||
|
|
|
|||
|
|
@ -599,7 +599,7 @@ func TestVulnerabilityListSupportsWheelAndPageNavigation(t *testing.T) {
|
|||
})
|
||||
}
|
||||
_, _, chatWidth, _ := model.layout()
|
||||
_, _, agentHeight := model.sidebarHeights()
|
||||
_, _, _, agentHeight := model.sidebarHeights()
|
||||
pageItems := model.vulnerabilityPageItems()
|
||||
|
||||
updated, _ := model.updateMouse(tea.MouseMsg{
|
||||
|
|
@ -923,7 +923,7 @@ func TestMainScrollbarsSupportClickAndDrag(t *testing.T) {
|
|||
model.viewport.SetContent(model.viewportContent)
|
||||
showSidebar, _, chatWidth, chatHeight := model.layout()
|
||||
viewerHeight := model.viewerHeight()
|
||||
_, vulnHeight, agentHeight := model.sidebarHeights()
|
||||
_, vulnHeight, _, agentHeight := model.sidebarHeights()
|
||||
if !showSidebar {
|
||||
t.Fatal("test requires sidebar")
|
||||
}
|
||||
|
|
@ -967,6 +967,61 @@ func TestMainScrollbarsSupportClickAndDrag(t *testing.T) {
|
|||
}
|
||||
}
|
||||
|
||||
func TestMcpRosterScrollsByKeyWheelAndScrollbar(t *testing.T) {
|
||||
model := New(nil)
|
||||
model.width, model.height = 150, 35
|
||||
model.ready = true
|
||||
conns := make([]protocol.Connection, 0, 12)
|
||||
for i := 0; i < 12; i++ {
|
||||
conns = append(conns, protocol.Connection{Name: fmt.Sprintf("conn-%02d", i), ToolCount: 2})
|
||||
}
|
||||
model.snapshot.Connections = conns
|
||||
|
||||
showSidebar, _, chatWidth, _ := model.layout()
|
||||
if !showSidebar {
|
||||
t.Fatal("test requires sidebar")
|
||||
}
|
||||
viewerHeight := model.viewerHeight()
|
||||
_, vulnHeight, mcpHeight, agentHeight := model.sidebarHeights()
|
||||
mcpTop := viewerHeight + agentHeight + vulnHeight
|
||||
bottom := model.clampMcpOffset(1 << 30)
|
||||
if bottom == 0 {
|
||||
t.Fatalf("a roster of %d should overflow the panel", len(conns))
|
||||
}
|
||||
|
||||
// Wheel over the panel focuses it and advances the window.
|
||||
updated, _ := model.updateMouse(tea.MouseMsg{
|
||||
X: chatWidth + 2, Y: mcpTop + 1, Button: tea.MouseButtonWheelDown,
|
||||
})
|
||||
model = updated.(Model)
|
||||
if model.focus != focusMcp || model.mcpOffset != 3 {
|
||||
t.Fatalf("wheel scroll did not focus and advance roster: focus=%v offset=%d", model.focus, model.mcpOffset)
|
||||
}
|
||||
|
||||
// Page down pins to the bottom; up steps back one.
|
||||
updated, _ = model.updateMain(tea.KeyMsg{Type: tea.KeyPgDown})
|
||||
model = updated.(Model)
|
||||
if model.mcpOffset != bottom {
|
||||
t.Fatalf("page down did not reach the roster bottom: offset=%d want=%d", model.mcpOffset, bottom)
|
||||
}
|
||||
updated, _ = model.updateMain(tea.KeyMsg{Type: tea.KeyUp})
|
||||
model = updated.(Model)
|
||||
if model.mcpOffset != bottom-1 {
|
||||
t.Fatalf("up did not step the roster back one: offset=%d want=%d", model.mcpOffset, bottom-1)
|
||||
}
|
||||
|
||||
// Clicking the scrollbar thumb captures it and moves the window.
|
||||
model.mcpOffset = 0
|
||||
updated, _ = model.updateMouse(tea.MouseMsg{
|
||||
X: model.width - 3, Y: mcpTop + mcpHeight - 2,
|
||||
Button: tea.MouseButtonLeft, Action: tea.MouseActionPress,
|
||||
})
|
||||
model = updated.(Model)
|
||||
if model.draggingScrollbar != scrollbarMcp || model.mcpOffset == 0 {
|
||||
t.Fatalf("mcp scrollbar click failed: drag=%v offset=%d", model.draggingScrollbar, model.mcpOffset)
|
||||
}
|
||||
}
|
||||
|
||||
func TestTerminalSnapshotWithoutAgentsDoesNotKeepLoading(t *testing.T) {
|
||||
tests := []struct {
|
||||
state string
|
||||
|
|
|
|||
|
|
@ -59,6 +59,14 @@ func (m Model) updateMain(key tea.KeyMsg) (tea.Model, tea.Cmd) {
|
|||
m.ensureVulnerabilityVisible()
|
||||
return m, nil
|
||||
}
|
||||
if m.focus == focusMcp && len(m.snapshot.Connections) > 0 {
|
||||
delta := 1
|
||||
if key.String() == "up" {
|
||||
delta = -1
|
||||
}
|
||||
m.mcpOffset = m.clampMcpOffset(m.mcpOffset + delta)
|
||||
return m, nil
|
||||
}
|
||||
case "enter", " ":
|
||||
if m.focus == focusVulnerabilities && len(m.snapshot.Vulnerabilities) > 0 {
|
||||
if key.String() == "enter" {
|
||||
|
|
@ -94,6 +102,10 @@ func (m Model) updateMain(key tea.KeyMsg) (tea.Model, tea.Cmd) {
|
|||
m.ensureVulnerabilityVisible()
|
||||
return m, nil
|
||||
}
|
||||
if m.focus == focusMcp && len(m.snapshot.Connections) > 0 {
|
||||
m.mcpOffset = m.clampMcpOffset(m.mcpOffset - m.mcpPageSize())
|
||||
return m, nil
|
||||
}
|
||||
m.focus = focusChat
|
||||
m.input.Blur()
|
||||
m.followOutput = false
|
||||
|
|
@ -105,6 +117,10 @@ func (m Model) updateMain(key tea.KeyMsg) (tea.Model, tea.Cmd) {
|
|||
m.ensureVulnerabilityVisible()
|
||||
return m, nil
|
||||
}
|
||||
if m.focus == focusMcp && len(m.snapshot.Connections) > 0 {
|
||||
m.mcpOffset = m.clampMcpOffset(m.mcpOffset + m.mcpPageSize())
|
||||
return m, nil
|
||||
}
|
||||
m.focus = focusChat
|
||||
m.input.Blur()
|
||||
m.viewport.HalfViewDown()
|
||||
|
|
@ -147,10 +163,10 @@ func (m Model) updateMouse(msg tea.MouseMsg) (tea.Model, tea.Cmd) {
|
|||
}
|
||||
showSidebar, _, chatWidth, chatHeight := m.layout()
|
||||
viewerHeight := m.viewerHeight()
|
||||
_, vulnHeight, agentHeight := m.sidebarHeights()
|
||||
_, vulnHeight, mcpHeight, agentHeight := m.sidebarHeights()
|
||||
x, y := msg.X, msg.Y
|
||||
if m.updateMainScrollbarMouse(
|
||||
msg, showSidebar, chatWidth, chatHeight, viewerHeight, agentHeight, vulnHeight,
|
||||
msg, showSidebar, chatWidth, chatHeight, viewerHeight, agentHeight, vulnHeight, mcpHeight,
|
||||
) {
|
||||
return m, nil
|
||||
}
|
||||
|
|
@ -196,6 +212,10 @@ func (m Model) updateMouse(msg tea.MouseMsg) (tea.Model, tea.Cmd) {
|
|||
m.input.Blur()
|
||||
m.vulnOffset = max(0, m.vulnOffset-3)
|
||||
m.keepVulnerabilitySelectionInWindow()
|
||||
case mcpHeight > 0 && y < viewerHeight+agentHeight+vulnHeight+mcpHeight:
|
||||
m.focus = focusMcp
|
||||
m.input.Blur()
|
||||
m.mcpOffset = m.clampMcpOffset(m.mcpOffset - 3)
|
||||
}
|
||||
return m, nil
|
||||
}
|
||||
|
|
@ -222,6 +242,10 @@ func (m Model) updateMouse(msg tea.MouseMsg) (tea.Model, tea.Cmd) {
|
|||
totalRows, _ := m.vulnerabilityScrollRows()
|
||||
m.vulnOffset = min(max(0, totalRows-m.vulnerabilityPageSize()), m.vulnOffset+3)
|
||||
m.keepVulnerabilitySelectionInWindow()
|
||||
case mcpHeight > 0 && y < viewerHeight+agentHeight+vulnHeight+mcpHeight:
|
||||
m.focus = focusMcp
|
||||
m.input.Blur()
|
||||
m.mcpOffset = m.clampMcpOffset(m.mcpOffset + 3)
|
||||
}
|
||||
return m, nil
|
||||
}
|
||||
|
|
@ -303,7 +327,7 @@ func (m Model) updateMouse(msg tea.MouseMsg) (tea.Model, tea.Cmd) {
|
|||
func (m *Model) updateMainScrollbarMouse(
|
||||
msg tea.MouseMsg,
|
||||
showSidebar bool,
|
||||
chatWidth, chatHeight, viewerHeight, agentHeight, vulnHeight int,
|
||||
chatWidth, chatHeight, viewerHeight, agentHeight, vulnHeight, mcpHeight int,
|
||||
) bool {
|
||||
if msg.Action == tea.MouseActionRelease {
|
||||
if m.draggingScrollbar == scrollbarNone {
|
||||
|
|
@ -313,18 +337,18 @@ func (m *Model) updateMainScrollbarMouse(
|
|||
return true
|
||||
}
|
||||
if msg.Action == tea.MouseActionMotion && m.draggingScrollbar != scrollbarNone {
|
||||
m.scrollFromMouse(m.draggingScrollbar, msg.Y, chatHeight, viewerHeight, agentHeight)
|
||||
m.scrollFromMouse(m.draggingScrollbar, msg.Y, chatHeight, viewerHeight, agentHeight, vulnHeight)
|
||||
return true
|
||||
}
|
||||
if msg.Action != tea.MouseActionPress || msg.Button != tea.MouseButtonLeft {
|
||||
return false
|
||||
}
|
||||
target := m.scrollbarAt(msg, showSidebar, chatWidth, chatHeight, viewerHeight, agentHeight, vulnHeight)
|
||||
target := m.scrollbarAt(msg, showSidebar, chatWidth, chatHeight, viewerHeight, agentHeight, vulnHeight, mcpHeight)
|
||||
if target == scrollbarNone {
|
||||
return false
|
||||
}
|
||||
m.draggingScrollbar = target
|
||||
m.scrollFromMouse(target, msg.Y, chatHeight, viewerHeight, agentHeight)
|
||||
m.scrollFromMouse(target, msg.Y, chatHeight, viewerHeight, agentHeight, vulnHeight)
|
||||
return true
|
||||
}
|
||||
|
||||
|
|
@ -341,8 +365,9 @@ func nearColumn(x, column int) bool {
|
|||
func (m Model) scrollbarAt(
|
||||
msg tea.MouseMsg,
|
||||
showSidebar bool,
|
||||
chatWidth, chatHeight, viewerHeight, agentHeight, vulnHeight int,
|
||||
chatWidth, chatHeight, viewerHeight, agentHeight, vulnHeight, mcpHeight int,
|
||||
) scrollbarTarget {
|
||||
mcpTop := viewerHeight + agentHeight + vulnHeight
|
||||
switch {
|
||||
case nearColumn(msg.X, chatWidth-2) && msg.Y >= 1 && msg.Y < chatHeight-1 &&
|
||||
m.viewport.TotalLineCount() > m.viewport.VisibleLineCount():
|
||||
|
|
@ -358,13 +383,20 @@ func (m Model) scrollbarAt(
|
|||
if totalRows > m.vulnerabilityPageSize() {
|
||||
return scrollbarFindings
|
||||
}
|
||||
// The roster scrolls below a fixed header, so its bar starts two rows into
|
||||
// the panel (border then header) rather than one.
|
||||
case showSidebar && mcpHeight > 0 && nearColumn(msg.X, m.width-3) &&
|
||||
msg.Y >= mcpTop+2 && msg.Y < mcpTop+mcpHeight-1:
|
||||
if len(m.snapshot.Connections) > m.mcpPageSize() {
|
||||
return scrollbarMcp
|
||||
}
|
||||
}
|
||||
return scrollbarNone
|
||||
}
|
||||
|
||||
func (m *Model) scrollFromMouse(
|
||||
target scrollbarTarget,
|
||||
y, chatHeight, viewerHeight, agentHeight int,
|
||||
y, chatHeight, viewerHeight, agentHeight, vulnHeight int,
|
||||
) {
|
||||
switch target {
|
||||
case scrollbarTrace:
|
||||
|
|
@ -390,6 +422,13 @@ func (m *Model) scrollFromMouse(
|
|||
// The offset is a row, so dragging moves the list continuously.
|
||||
m.vulnOffset = scrollbarOffset(y-viewerHeight-agentHeight-1, height, totalRows, height)
|
||||
m.keepVulnerabilitySelectionInWindow()
|
||||
case scrollbarMcp:
|
||||
height := m.mcpPageSize()
|
||||
total := len(m.snapshot.Connections)
|
||||
m.focus = focusMcp
|
||||
m.input.Blur()
|
||||
// The bar starts two rows into the panel (border then the fixed header).
|
||||
m.mcpOffset = scrollbarOffset(y-viewerHeight-agentHeight-vulnHeight-2, height, total, height)
|
||||
}
|
||||
}
|
||||
|
||||
|
|
@ -545,6 +584,9 @@ func (m *Model) cycleFocus(delta int) {
|
|||
if len(m.snapshot.Vulnerabilities) > 0 {
|
||||
available = append(available, focusVulnerabilities)
|
||||
}
|
||||
if len(m.snapshot.Connections) > 0 {
|
||||
available = append(available, focusMcp)
|
||||
}
|
||||
}
|
||||
idx := 0
|
||||
for i, focus := range available {
|
||||
|
|
|
|||
|
|
@ -507,7 +507,7 @@ func (m Model) mainView() string {
|
|||
func (m Model) sidebarView(width, height int) string {
|
||||
// Stats box height fits its content (auto, max 15); vulns panel max-height 12.
|
||||
statsBody := m.statsView()
|
||||
statsHeight, vulnHeight, agentHeight := m.sidebarHeights()
|
||||
statsHeight, vulnHeight, mcpHeight, agentHeight := m.sidebarHeights()
|
||||
agentBorder := dark
|
||||
if m.focus == focusAgents {
|
||||
agentBorder = green
|
||||
|
|
@ -546,11 +546,19 @@ func (m Model) sidebarView(width, height int) string {
|
|||
)
|
||||
parts = append(parts, lipgloss.NewStyle().Width(width-2).Height(vulnRows).Border(lipgloss.RoundedBorder()).BorderForeground(vulnBorder).Padding(0, 1).Render(findings))
|
||||
}
|
||||
if mcpHeight > 0 {
|
||||
mcpBorder := dark
|
||||
if m.focus == focusMcp {
|
||||
mcpBorder = green
|
||||
}
|
||||
mcpRows := max(1, mcpHeight-2)
|
||||
parts = append(parts, lipgloss.NewStyle().Width(width-2).Height(mcpRows).Border(lipgloss.RoundedBorder()).BorderForeground(mcpBorder).Padding(0, 1).Render(m.mcpConnectionsView(width-4, mcpRows)))
|
||||
}
|
||||
parts = append(parts, lipgloss.NewStyle().Width(width-2).Height(statsHeight-2).Border(lipgloss.RoundedBorder()).BorderForeground(dark).Padding(0, 1).Render(statsBody))
|
||||
return strings.Join(parts, "\n")
|
||||
}
|
||||
|
||||
func (m Model) sidebarHeights() (statsHeight, vulnHeight, agentHeight int) {
|
||||
func (m Model) sidebarHeights() (statsHeight, vulnHeight, mcpHeight, agentHeight int) {
|
||||
// Measure the stats panel the way its box will render it: a long model name
|
||||
// wraps inside the sidebar, and counting only its newlines would size the
|
||||
// box short and push the whole frame past the bottom of the terminal.
|
||||
|
|
@ -559,7 +567,13 @@ func (m Model) sidebarHeights() (statsHeight, vulnHeight, agentHeight int) {
|
|||
if len(m.snapshot.Vulnerabilities) > 0 {
|
||||
vulnHeight = min(12, len(m.vulnerabilityRows(m.vulnerabilityListWidth()))+2)
|
||||
}
|
||||
agentHeight = max(3, m.height-m.viewerHeight()-statsHeight-vulnHeight)
|
||||
// One header line + one line per connection + the box border (2). Capped so a
|
||||
// long roster cannot crowd out the agent tree; a roster past the cap scrolls
|
||||
// inside the panel. Absent entirely when the run has no MCP connections.
|
||||
if len(m.snapshot.Connections) > 0 {
|
||||
mcpHeight = min(9, len(m.snapshot.Connections)+3)
|
||||
}
|
||||
agentHeight = max(3, m.height-m.viewerHeight()-statsHeight-vulnHeight-mcpHeight)
|
||||
return
|
||||
}
|
||||
|
||||
|
|
@ -628,6 +642,111 @@ func (m Model) statsView() string {
|
|||
return b.String()
|
||||
}
|
||||
|
||||
// mcpConnectionsView renders the sidebar MCP panel: a header carrying the total
|
||||
// connection count, then one row per connection with a status glyph and its tool
|
||||
// count (or "offline").
|
||||
// - a solid green dot marks an attached, idle connection;
|
||||
// - a green cycling quarter-circle (◐ ◓ ◑ ◒) marks a call running against it;
|
||||
// - a red dot plus "offline" marks a connection whose live session has died.
|
||||
//
|
||||
// The header stays fixed while the roster below it scrolls: when there are more
|
||||
// connections than the panel can show, the visible window is chosen by
|
||||
// m.mcpOffset and withVerticalScrollbar draws a thumb in the reserved last
|
||||
// column, exactly as the agent tree and findings list scroll.
|
||||
//
|
||||
// "In use" is derived from the connection-tagged tool-call events in the stream,
|
||||
// not carried on the connection roster, so a call in flight shows motion without
|
||||
// any extra backend signal. The quarter-circle rides the shared sweepFrame tick.
|
||||
func (m Model) mcpConnectionsView(width, rows int) string {
|
||||
conns := m.snapshot.Connections
|
||||
header := truncate(lipgloss.NewStyle().Foreground(dim).Render(
|
||||
fmt.Sprintf("MCP Connections (%d)", len(conns))), width)
|
||||
bodyRows := max(0, rows-1)
|
||||
if bodyRows == 0 {
|
||||
return header
|
||||
}
|
||||
inUse := m.mcpInUse()
|
||||
frames := []rune{'◐', '◓', '◑', '◒'}
|
||||
// Reserve the scrollbar column whether or not the bar is showing, so the
|
||||
// roster does not shift sideways as it grows past the panel.
|
||||
rosterWidth := max(1, width-1)
|
||||
start := windowStart(m.mcpOffset, len(conns), bodyRows)
|
||||
end := min(len(conns), start+bodyRows)
|
||||
lines := make([]string, 0, max(0, end-start))
|
||||
for i := start; i < end; i++ {
|
||||
conn := conns[i]
|
||||
var glyph, right string
|
||||
switch {
|
||||
case conn.Dead:
|
||||
glyph = lipgloss.NewStyle().Foreground(red).Render("●")
|
||||
right = lipgloss.NewStyle().Foreground(red).Render("offline")
|
||||
case inUse[conn.Name]:
|
||||
glyph = lipgloss.NewStyle().Foreground(green).Render(string(frames[m.sweepFrame%len(frames)]))
|
||||
right = lipgloss.NewStyle().Foreground(dim).Render(toolsLabel(conn.ToolCount))
|
||||
default:
|
||||
glyph = lipgloss.NewStyle().Foreground(green).Render("●")
|
||||
right = lipgloss.NewStyle().Foreground(dim).Render(toolsLabel(conn.ToolCount))
|
||||
}
|
||||
rightWidth := lipgloss.Width(right)
|
||||
name := truncate(lipgloss.NewStyle().Foreground(textColor).Render(conn.Name), max(1, rosterWidth-2-rightWidth-1))
|
||||
gap := max(1, rosterWidth-2-lipgloss.Width(name)-rightWidth)
|
||||
lines = append(lines, glyph+" "+name+strings.Repeat(" ", gap)+right)
|
||||
}
|
||||
roster := withVerticalScrollbar(
|
||||
strings.Join(lines, "\n"),
|
||||
width,
|
||||
bodyRows,
|
||||
len(conns),
|
||||
bodyRows,
|
||||
m.mcpOffset,
|
||||
m.scrollbarThumb(scrollbarMcp),
|
||||
)
|
||||
return header + "\n" + roster
|
||||
}
|
||||
|
||||
// mcpPageSize is how many connection rows the roster shows at once, below its
|
||||
// fixed header line.
|
||||
func (m Model) mcpPageSize() int {
|
||||
_, _, mcpHeight, _ := m.sidebarHeights()
|
||||
// mcpHeight = 2 (border) + header (1) + roster rows.
|
||||
return max(1, mcpHeight-3)
|
||||
}
|
||||
|
||||
// clampMcpOffset keeps the roster offset within the range that still shows a
|
||||
// full page of connections at the bottom.
|
||||
func (m Model) clampMcpOffset(offset int) int {
|
||||
return min(max(0, offset), max(0, len(m.snapshot.Connections)-m.mcpPageSize()))
|
||||
}
|
||||
|
||||
// mcpInUse is the set of MCP connections with a tool call currently running,
|
||||
// read off the connection-tagged tool events the model already holds. Each MCP
|
||||
// dispatch event carries the connection name (mcp_connection) and a status that
|
||||
// moves running -> completed as its own event is upserted, so a connection is
|
||||
// "in use" exactly while one of its events is still running.
|
||||
func (m Model) mcpInUse() map[string]bool {
|
||||
inUse := map[string]bool{}
|
||||
for _, event := range m.snapshot.Events {
|
||||
if event.Type != "tool" {
|
||||
continue
|
||||
}
|
||||
connection := render.StringValue(event.Data["mcp_connection"])
|
||||
if connection == "" {
|
||||
continue
|
||||
}
|
||||
if render.StringValue(event.Data["status"]) == "running" {
|
||||
inUse[connection] = true
|
||||
}
|
||||
}
|
||||
return inUse
|
||||
}
|
||||
|
||||
func toolsLabel(count int) string {
|
||||
if count == 1 {
|
||||
return "1 tool"
|
||||
}
|
||||
return fmt.Sprintf("%d tools", count)
|
||||
}
|
||||
|
||||
func numberValue(value any) int64 {
|
||||
switch v := value.(type) {
|
||||
case float64:
|
||||
|
|
|
|||
|
|
@ -134,7 +134,7 @@ func clampVulnerabilityOffset(offset, total, height int) int {
|
|||
}
|
||||
|
||||
func (m Model) vulnerabilityPageSize() int {
|
||||
_, vulnHeight, _ := m.sidebarHeights()
|
||||
_, vulnHeight, _, _ := m.sidebarHeights()
|
||||
return max(1, vulnHeight-2)
|
||||
}
|
||||
|
||||
|
|
|
|||
|
|
@ -32,6 +32,17 @@ type Agent struct {
|
|||
ErrorMessage string `json:"error_message"`
|
||||
}
|
||||
|
||||
// Connection is one MCP connection the run may reach, as the backend projects
|
||||
// it for the sidebar's MCP panel. Non-secret by construction: only the display
|
||||
// name, how many tools the connection offers, and whether its live session has
|
||||
// died (its reconnect-retry gave up). "In use" is not carried here; the client
|
||||
// derives it from the connection-tagged tool-call events in the event stream.
|
||||
type Connection struct {
|
||||
Name string `json:"name"`
|
||||
ToolCount int `json:"tool_count"`
|
||||
Dead bool `json:"dead"`
|
||||
}
|
||||
|
||||
type Event struct {
|
||||
ID string `json:"id"`
|
||||
Type string `json:"type"`
|
||||
|
|
@ -68,6 +79,7 @@ type Snapshot struct {
|
|||
Vulnerabilities []map[string]any `json:"-"`
|
||||
Usage map[string]any `json:"usage"`
|
||||
Subscription bool `json:"subscription"`
|
||||
Connections []Connection `json:"connections"`
|
||||
ViewerStatus string `json:"viewer_status"`
|
||||
ViewerURL *string `json:"viewer_url"`
|
||||
Error *string `json:"error"`
|
||||
|
|
|
|||
|
|
@ -45,3 +45,51 @@ func renderMcpInspect(connection, status string) string {
|
|||
b.WriteString(style.Render(icon))
|
||||
return b.String()
|
||||
}
|
||||
|
||||
// renderMcpList renders list_mcps: the inventory of connections the run may
|
||||
// reach, not a call to any of them, so no connection leads and the event
|
||||
// carries no connection tag. Unlike the other MCP results, the names are worth
|
||||
// showing: Strix assembled them itself from the run's registered connections,
|
||||
// so they are short and never an outside server's payload.
|
||||
func renderMcpList(result any, status string) string {
|
||||
var b strings.Builder
|
||||
b.WriteString(mcpIcon + Dim().Render("Listing MCP servers") + "\n")
|
||||
for _, conn := range mcpConnectionEntries(result) {
|
||||
b.WriteString(" " + Col(Slate).Render(conn.name))
|
||||
if conn.dead {
|
||||
b.WriteString(Dim().Render(" · ") + Col(Red).Render("offline"))
|
||||
}
|
||||
b.WriteString("\n")
|
||||
}
|
||||
icon, style := statusIcon(status)
|
||||
b.WriteString(style.Render(icon))
|
||||
return b.String()
|
||||
}
|
||||
|
||||
// mcpListEntry is one connection read out of a list_mcps result: its display
|
||||
// name and whether its live session has died.
|
||||
type mcpListEntry struct {
|
||||
name string
|
||||
dead bool
|
||||
}
|
||||
|
||||
// mcpConnectionEntries reads the connections out of a list_mcps result, which is
|
||||
// {"connections": [{"name": ..., "dead": ...}, ...]}. Anything else (still
|
||||
// running, or a result bounded down to a string) yields no entries, and the
|
||||
// header plus status stand alone.
|
||||
func mcpConnectionEntries(result any) []mcpListEntry {
|
||||
resultMap, _ := result.(map[string]any)
|
||||
connections, _ := resultMap["connections"].([]any)
|
||||
var entries []mcpListEntry
|
||||
for _, raw := range connections {
|
||||
entry, ok := raw.(map[string]any)
|
||||
if !ok {
|
||||
continue
|
||||
}
|
||||
if name := strings.TrimSpace(StringValue(entry["name"])); name != "" {
|
||||
dead, _ := entry["dead"].(bool)
|
||||
entries = append(entries, mcpListEntry{name: name, dead: dead})
|
||||
}
|
||||
}
|
||||
return entries
|
||||
}
|
||||
|
|
|
|||
|
|
@ -70,6 +70,10 @@ func Tool(data map[string]any) string {
|
|||
}
|
||||
|
||||
switch name {
|
||||
// list_mcps inventories every connection rather than touching one, so it is
|
||||
// the one MCP tool with no connection tag and routes by name like a built-in.
|
||||
case "list_mcps":
|
||||
return renderMcpList(result, status)
|
||||
case "exec_command":
|
||||
return renderExecCommand(args, result, status)
|
||||
case "write_stdin":
|
||||
|
|
|
|||
|
|
@ -268,6 +268,24 @@ func TestMcpDescribeInspectsConnection(t *testing.T) {
|
|||
}
|
||||
}
|
||||
|
||||
func TestMcpListMarksDeadConnectionsOffline(t *testing.T) {
|
||||
// list_mcps carries a per-connection dead flag; a dead connection reads as
|
||||
// offline in the inventory while a live one shows normally.
|
||||
result := map[string]any{
|
||||
"connections": []any{
|
||||
map[string]any{"name": "supabase", "tool_count": float64(3), "dead": false},
|
||||
map[string]any{"name": "vercel", "tool_count": float64(1), "dead": true},
|
||||
},
|
||||
}
|
||||
data := tool("list_mcps", nil, result, "completed")
|
||||
|
||||
out := ansi.Strip(Tool(data))
|
||||
requireContains(t, out, "Listing MCP servers", "supabase", "vercel", "offline")
|
||||
if strings.Count(out, "offline") != 1 {
|
||||
t.Fatalf("only the dead connection should read offline:\n%s", out)
|
||||
}
|
||||
}
|
||||
|
||||
func TestCollapseToolShellPreviewAndExpand(t *testing.T) {
|
||||
lines := make([]string, 16)
|
||||
for i := range lines {
|
||||
|
|
|
|||
|
|
@ -185,6 +185,7 @@ class GoTuiRuntime:
|
|||
max_turns=self.args.max_turns,
|
||||
max_budget_usd=self.args.max_budget_usd,
|
||||
event_sink=self.capture_event,
|
||||
mcp_status_sink=self.capture_mcp_status,
|
||||
)
|
||||
await self._sync_agent_state()
|
||||
if self.controller.scan_state == "running":
|
||||
|
|
@ -210,6 +211,15 @@ class GoTuiRuntime:
|
|||
self.live_view.ingest_sdk_event(agent_id, event)
|
||||
self.controller.notify_changed()
|
||||
|
||||
def capture_mcp_status(self, roster: list[dict[str, Any]]) -> None:
|
||||
"""Receive the engine's MCP connection roster and hand it to the controller.
|
||||
|
||||
Runs on the scan's event loop (called from the runner at establishment
|
||||
and from a session's on-dead callback), the same loop that drives
|
||||
``capture_event``, so updating the controller and repainting here is
|
||||
safe. The controller renders it as the sidebar MCP connections panel."""
|
||||
self.controller.set_mcp_connections(roster)
|
||||
|
||||
async def _sync_agent_state(self) -> bool:
|
||||
parent_of, statuses, names, errors = await self.coordinator.graph_snapshot()
|
||||
changed = False
|
||||
|
|
|
|||
|
|
@ -30,6 +30,7 @@ import {
|
|||
fetchTranscript,
|
||||
fetchVulnerabilities,
|
||||
forgetAuth,
|
||||
parseMcpConnectionStatus,
|
||||
type AuthStatus,
|
||||
type LoadedRun,
|
||||
type RunsPayload,
|
||||
|
|
@ -169,6 +170,27 @@ export default function App() {
|
|||
const agentCount = run?.transcript.agents.length ?? 0;
|
||||
const verified = auth?.verified === true;
|
||||
|
||||
// The run's persisted MCP roster (from run.json via /api/run), plus the set of
|
||||
// connections with a tool call currently in flight. "In use" is derived here
|
||||
// from the connection-tagged tool events rather than carried on the roster:
|
||||
// an MCP dispatch event carries its connection name and a status that moves
|
||||
// running -> completed, so a connection is in use while one of its events is
|
||||
// still running. This mirrors the terminal UI's MCP panel exactly.
|
||||
const mcpConnections = useMemo(
|
||||
() => (run ? parseMcpConnectionStatus(run.raw) : []),
|
||||
[run]
|
||||
);
|
||||
const mcpInUse = useMemo(() => {
|
||||
const inUse = new Set<string>();
|
||||
for (const event of run?.transcript.events ?? []) {
|
||||
if (event.type !== "tool") continue;
|
||||
const connection = event.data?.mcp_connection;
|
||||
if (typeof connection !== "string" || !connection) continue;
|
||||
if (event.data?.status === "running") inUse.add(connection);
|
||||
}
|
||||
return inUse;
|
||||
}, [run]);
|
||||
|
||||
// Per-run guard for the default view: land on Agents while a scan is live,
|
||||
// Overview once it finishes. Applied at most once per run and never once the
|
||||
// user has navigated manually (userSetView flips the guard).
|
||||
|
|
@ -251,6 +273,8 @@ export default function App() {
|
|||
}}
|
||||
issuesCount={run?.vulnerabilities.length ?? 0}
|
||||
agentCount={agentCount}
|
||||
mcpConnections={mcpConnections}
|
||||
mcpInUse={mcpInUse}
|
||||
runCount={runs?.count ?? 0}
|
||||
finished={run?.finished ?? false}
|
||||
verified={verified}
|
||||
|
|
|
|||
|
|
@ -14,6 +14,7 @@ import { IoChatbubblesOutline } from "react-icons/io5";
|
|||
import { cn } from "@/lib/utils";
|
||||
import { ctaUrl, trackCta } from "@/lib/cta";
|
||||
import { UpgradeModal } from "@/components/UpgradeModal";
|
||||
import type { McpConnectionStatus } from "@/data/serverSource";
|
||||
import type { View } from "@/App";
|
||||
|
||||
/**
|
||||
|
|
@ -37,6 +38,8 @@ interface SidebarProps {
|
|||
onSelectView: (view: View) => void;
|
||||
issuesCount: number;
|
||||
agentCount: number;
|
||||
mcpConnections: McpConnectionStatus[];
|
||||
mcpInUse: Set<string>;
|
||||
runCount: number;
|
||||
finished: boolean;
|
||||
verified: boolean;
|
||||
|
|
@ -61,6 +64,8 @@ export default function Sidebar({
|
|||
onSelectView,
|
||||
issuesCount,
|
||||
agentCount,
|
||||
mcpConnections,
|
||||
mcpInUse,
|
||||
runCount,
|
||||
finished,
|
||||
verified,
|
||||
|
|
@ -246,6 +251,9 @@ export default function Sidebar({
|
|||
onClick={() => onSelectView("agents")}
|
||||
/>
|
||||
)}
|
||||
{mcpConnections.length > 0 && (
|
||||
<McpConnectionsPanel connections={mcpConnections} inUse={mcpInUse} />
|
||||
)}
|
||||
<NavItem
|
||||
icon={<History className="h-4 w-4" />}
|
||||
label="Past runs"
|
||||
|
|
@ -421,6 +429,84 @@ function NavItem({ icon, label, active, onClick, count }: NavItemProps) {
|
|||
);
|
||||
}
|
||||
|
||||
// The quarter-circle sweep frames the terminal UI cycles for an in-use
|
||||
// connection, and the sub-second tick that advances them.
|
||||
const SWEEP_FRAMES = ["◐", "◓", "◑", "◒"] as const;
|
||||
const SWEEP_MS = 220;
|
||||
|
||||
/**
|
||||
* The MCP connections panel: a compact roster of the run's connected MCP
|
||||
* servers, matching the terminal UI's sidebar panel. A header carries the
|
||||
* total count; each row shows a status glyph, the connection name, and its
|
||||
* tool count (or "offline"):
|
||||
* - solid green dot: attached and idle;
|
||||
* - green cycling quarter-circle (◐◓◑◒): a tool call is running against it;
|
||||
* - red dot + "offline": the connection's live session has died.
|
||||
*
|
||||
* "In use" is derived by the caller from the connection-tagged tool events, not
|
||||
* carried on the roster, so a call in flight shows motion with no extra signal.
|
||||
* The roster scrolls within a bounded height so a long list never blows out the
|
||||
* rail, mirroring how the nav above it scrolls.
|
||||
*/
|
||||
function McpConnectionsPanel({
|
||||
connections,
|
||||
inUse,
|
||||
}: {
|
||||
connections: McpConnectionStatus[];
|
||||
inUse: Set<string>;
|
||||
}) {
|
||||
const anyInUse = connections.some((c) => !c.dead && inUse.has(c.name));
|
||||
const [frame, setFrame] = useState(0);
|
||||
|
||||
// Advance the sweep only while at least one connection is in use, so an idle
|
||||
// panel does no work.
|
||||
useEffect(() => {
|
||||
if (!anyInUse) return;
|
||||
const id = setInterval(() => setFrame((f) => (f + 1) % SWEEP_FRAMES.length), SWEEP_MS);
|
||||
return () => clearInterval(id);
|
||||
}, [anyInUse]);
|
||||
|
||||
return (
|
||||
<div className="mt-1">
|
||||
<div className="flex h-7 items-center px-2 text-[11px] font-medium text-[#666]">
|
||||
MCP Connections ({connections.length})
|
||||
</div>
|
||||
<div className="max-h-48 overflow-y-auto overflow-x-clip scrollbar-thin">
|
||||
{connections.map((conn) => {
|
||||
const busy = !conn.dead && inUse.has(conn.name);
|
||||
return (
|
||||
<div
|
||||
key={conn.name}
|
||||
className="flex h-7 items-center gap-2 rounded-md px-2"
|
||||
title={conn.provider ? `${conn.name} · ${conn.provider}` : conn.name}
|
||||
>
|
||||
<span
|
||||
className={cn(
|
||||
"w-3 flex-none text-center text-[11px] leading-none",
|
||||
conn.dead ? "text-red-400" : "text-emerald-400"
|
||||
)}
|
||||
aria-hidden="true"
|
||||
>
|
||||
{conn.dead ? "●" : busy ? SWEEP_FRAMES[frame] : "●"}
|
||||
</span>
|
||||
<span className="min-w-0 flex-1 truncate text-[13px] font-medium text-[#ededed]">
|
||||
{conn.name}
|
||||
</span>
|
||||
{conn.dead ? (
|
||||
<span className="flex-none text-[11px] text-red-400">offline</span>
|
||||
) : (
|
||||
<span className="flex-none text-[11px] tabular-nums text-[#666]">
|
||||
{conn.toolCount} {conn.toolCount === 1 ? "tool" : "tools"}
|
||||
</span>
|
||||
)}
|
||||
</div>
|
||||
);
|
||||
})}
|
||||
</div>
|
||||
</div>
|
||||
);
|
||||
}
|
||||
|
||||
// Overview icon: a dashboard grid glyph (16x16 viewBox).
|
||||
function ProjectsIcon() {
|
||||
return (
|
||||
|
|
|
|||
|
|
@ -14,6 +14,10 @@ import type { ToolRendererProps } from "@/types/events";
|
|||
* never as markdown, since it came from a server outside Strix.
|
||||
*
|
||||
* The full result is still in the run's event data on disk either way.
|
||||
*
|
||||
* list_mcps is the other exception: its result is the engine's own inventory of
|
||||
* the run's connections (names and tool counts), short and assembled by Strix
|
||||
* rather than returned by an outside server, so it is shown inline.
|
||||
*/
|
||||
|
||||
/** Arguments one line each, as the terminal prints them. */
|
||||
|
|
@ -25,6 +29,52 @@ function argLines(args: unknown): string[] {
|
|||
});
|
||||
}
|
||||
|
||||
/** One connection out of a list_mcps inventory. */
|
||||
interface McpListingEntry {
|
||||
name: string;
|
||||
toolCount: number | null;
|
||||
dead: boolean;
|
||||
}
|
||||
|
||||
/**
|
||||
* The connections out of a list_mcps result, which is
|
||||
* `{"connections": [{id, name, description, tool_count}, ...]}`, sometimes
|
||||
* arriving JSON-encoded as a string. Anything else yields an empty list and the
|
||||
* row shows just the header and status. Unlike other MCP results this one is
|
||||
* safe to show: the engine assembled it from the run's own registered
|
||||
* connections, so it is short and never an outside server's payload. It still
|
||||
* renders as inert text.
|
||||
*/
|
||||
function listingEntries(result: unknown): McpListingEntry[] {
|
||||
let value = result;
|
||||
if (typeof value === "string") {
|
||||
try {
|
||||
value = JSON.parse(value);
|
||||
} catch {
|
||||
return [];
|
||||
}
|
||||
}
|
||||
const connections =
|
||||
value && typeof value === "object" && !Array.isArray(value)
|
||||
? (value as Record<string, unknown>).connections
|
||||
: null;
|
||||
if (!Array.isArray(connections)) return [];
|
||||
return connections.flatMap((entry) => {
|
||||
if (!entry || typeof entry !== "object" || Array.isArray(entry)) return [];
|
||||
const record = entry as Record<string, unknown>;
|
||||
const name =
|
||||
typeof record.name === "string" && record.name.trim()
|
||||
? record.name.trim()
|
||||
: typeof record.id === "string"
|
||||
? record.id.trim()
|
||||
: "";
|
||||
if (!name) return [];
|
||||
const toolCount = typeof record.tool_count === "number" ? record.tool_count : null;
|
||||
const dead = record.dead === true;
|
||||
return [{ name, toolCount, dead }];
|
||||
});
|
||||
}
|
||||
|
||||
const MAX_ERROR_CHARS = 600;
|
||||
|
||||
function errorText(result: unknown): string | null {
|
||||
|
|
@ -50,11 +100,17 @@ export default function McpRenderer({
|
|||
// describe_mcp inspects a connection's catalog rather than calling a tool on
|
||||
// it, so the connection is the subject and there is no underlying tool.
|
||||
const inspecting = toolName === "describe_mcp";
|
||||
// list_mcps inventories every connection rather than touching one, so it
|
||||
// carries no connection at all and is routed here by name instead.
|
||||
const listing = toolName === "list_mcps";
|
||||
const entries = listing ? listingEntries(result) : [];
|
||||
|
||||
return (
|
||||
<div>
|
||||
<div className="flex items-center gap-2 flex-wrap">
|
||||
{inspecting ? (
|
||||
{listing ? (
|
||||
<span className="text-[13px] text-[#555]">Listing connected MCP servers</span>
|
||||
) : inspecting ? (
|
||||
<>
|
||||
<span className="text-[13px] text-[#555]">Inspecting MCP server</span>
|
||||
{mcpConnection && (
|
||||
|
|
@ -82,6 +138,26 @@ export default function McpRenderer({
|
|||
</div>
|
||||
)}
|
||||
|
||||
{entries.length > 0 && (
|
||||
<div className="mt-1 font-mono text-[13px] leading-relaxed">
|
||||
{entries.map((entry) => (
|
||||
<div key={entry.name} className={`break-all${entry.dead ? " opacity-50" : ""}`}>
|
||||
<span className="text-teal-300">{entry.name}</span>
|
||||
{entry.dead ? (
|
||||
<span className="text-red-400/80"> · offline</span>
|
||||
) : (
|
||||
entry.toolCount !== null && (
|
||||
<span className="text-[#555]">
|
||||
{" "}
|
||||
· {entry.toolCount} {entry.toolCount === 1 ? "tool" : "tools"}
|
||||
</span>
|
||||
)
|
||||
)}
|
||||
</div>
|
||||
))}
|
||||
</div>
|
||||
)}
|
||||
|
||||
<div className="mt-1 text-[13px]">
|
||||
{status === "running" && <span className="text-[#666]">Running</span>}
|
||||
{status === "completed" && <span className="text-emerald-400/80">✓ Done</span>}
|
||||
|
|
|
|||
|
|
@ -93,7 +93,9 @@ const CATEGORY_META: Record<ToolCategory, CategoryMeta> = {
|
|||
threatModel: { renderer: ThreatModelRenderer, icon: Crosshair, color: "text-blue-400", match: /threat_model/ },
|
||||
telemetry: { renderer: FallbackRenderer, icon: Wrench, color: "text-[#555]" },
|
||||
// Tools from the user's own MCP servers. Resolved from the connection on the
|
||||
// event rather than from a tool name, so this family has no names below.
|
||||
// event rather than from a tool name — except list_mcps, the engine's
|
||||
// inventory of every connection, which touches none and so carries no
|
||||
// connection to resolve from; it is the family's one name below.
|
||||
mcp: { renderer: McpRenderer, icon: Plug, color: "text-teal-400" },
|
||||
};
|
||||
|
||||
|
|
@ -128,7 +130,7 @@ const CATEGORY_TOOLS: Record<ToolCategory, readonly string[]> = {
|
|||
// Per-target threat model, shared across the agent tree
|
||||
threatModel: ["get_threat_model", "save_threat_model", "amend_threat_model"],
|
||||
telemetry: ["sandbox_error_details", "llm_error_details"],
|
||||
mcp: [],
|
||||
mcp: ["list_mcps"],
|
||||
};
|
||||
|
||||
/** Reverse index (tool name → family), built once from CATEGORY_TOOLS. */
|
||||
|
|
|
|||
|
|
@ -39,6 +39,39 @@ export interface Transcript {
|
|||
events: TranscriptEvent[];
|
||||
}
|
||||
|
||||
/**
|
||||
* One MCP connection's non-secret status, as persisted to run.json by the
|
||||
* engine under `mcp_connection_status` and surfaced verbatim by GET /api/run.
|
||||
* Only name / provider / tool_count / dead ride here; never config, url, or
|
||||
* token. `dead` means the connection's live session gave up reconnecting.
|
||||
*/
|
||||
export interface McpConnectionStatus {
|
||||
name: string;
|
||||
provider: string | null;
|
||||
toolCount: number;
|
||||
dead: boolean;
|
||||
}
|
||||
|
||||
/**
|
||||
* Read the MCP connection roster out of a raw run record. Tolerates the field
|
||||
* being absent (older runs, or a run with no MCP) and any malformed entry,
|
||||
* yielding an empty list rather than throwing.
|
||||
*/
|
||||
export function parseMcpConnectionStatus(raw: Record<string, unknown>): McpConnectionStatus[] {
|
||||
const list = raw?.mcp_connection_status;
|
||||
if (!Array.isArray(list)) return [];
|
||||
return list.flatMap((entry) => {
|
||||
if (!entry || typeof entry !== "object" || Array.isArray(entry)) return [];
|
||||
const record = entry as Record<string, unknown>;
|
||||
const name = typeof record.name === "string" ? record.name.trim() : "";
|
||||
if (!name) return [];
|
||||
const provider = typeof record.provider === "string" && record.provider.trim() ? record.provider.trim() : null;
|
||||
const toolCount = typeof record.tool_count === "number" ? record.tool_count : 0;
|
||||
const dead = record.dead === true;
|
||||
return [{ name, provider, toolCount, dead }];
|
||||
});
|
||||
}
|
||||
|
||||
export interface LoadedRun {
|
||||
summary: ParsedRunSummary;
|
||||
/** Whole raw run record (for llm_usage, targets_info details, etc.). */
|
||||
|
|
|
|||
File diff suppressed because one or more lines are too long
File diff suppressed because one or more lines are too long
10
strix/interface/viewer/static/assets/index-qwPOPAGC.css
Normal file
10
strix/interface/viewer/static/assets/index-qwPOPAGC.css
Normal file
File diff suppressed because one or more lines are too long
|
|
@ -6,8 +6,8 @@
|
|||
<meta name="viewport" content="width=device-width, initial-scale=1.0" />
|
||||
<meta name="color-scheme" content="dark" />
|
||||
<title>Strix Results</title>
|
||||
<script type="module" crossorigin src="./assets/index-CYf9nnT3.js"></script>
|
||||
<link rel="stylesheet" crossorigin href="./assets/index-D0453ODW.css">
|
||||
<script type="module" crossorigin src="./assets/index-Bpn8GiSb.js"></script>
|
||||
<link rel="stylesheet" crossorigin href="./assets/index-qwPOPAGC.css">
|
||||
</head>
|
||||
<body>
|
||||
<div id="root"></div>
|
||||
|
|
|
|||
|
|
@ -416,6 +416,22 @@ class ReportState:
|
|||
self.run_record["mcp_connections"] = names
|
||||
self.save_run_data()
|
||||
|
||||
def record_mcp_connection_status(self, status: list[dict[str, Any]]) -> None:
|
||||
"""Persist the run's non-secret MCP connection status roster.
|
||||
|
||||
``status`` is one entry per connection carrying only ``name``,
|
||||
``provider``, ``tool_count``, and ``dead`` (no config, url, token, or
|
||||
auth). Saved as soon as the run connects and rewritten each time a
|
||||
connection dies, so the viewer, which rebuilds its display by re-reading
|
||||
the run's files from disk, can show a live connections panel and health
|
||||
without any in-memory event sink. Kept separate from the
|
||||
``mcp_connections`` name list so neither field repurposes the other.
|
||||
"""
|
||||
if self.run_record.get("mcp_connection_status") == status:
|
||||
return
|
||||
self.run_record["mcp_connection_status"] = status
|
||||
self.save_run_data()
|
||||
|
||||
def set_scan_config(self, config: dict[str, Any]) -> None:
|
||||
self.scan_config = config
|
||||
self.run_record["status"] = "running"
|
||||
|
|
|
|||
|
|
@ -23,10 +23,12 @@ from strix.tools.mcp.registry import (
|
|||
McpCallInfo,
|
||||
McpConnectionEntry,
|
||||
McpConnectionRequest,
|
||||
McpConnectionStatus,
|
||||
McpConnectionSummary,
|
||||
McpRegistry,
|
||||
resolve_mcp_call,
|
||||
)
|
||||
from strix.tools.mcp.session import McpConnectionUnavailableError, SupervisedMcpSession
|
||||
|
||||
|
||||
__all__ = [
|
||||
|
|
@ -41,8 +43,11 @@ __all__ = [
|
|||
"McpConnectionConfig",
|
||||
"McpConnectionEntry",
|
||||
"McpConnectionRequest",
|
||||
"McpConnectionStatus",
|
||||
"McpConnectionSummary",
|
||||
"McpConnectionUnavailableError",
|
||||
"McpRegistry",
|
||||
"SupervisedMcpSession",
|
||||
"attach_mcp_requests",
|
||||
"call_mcp",
|
||||
"connect_mcp_servers",
|
||||
|
|
|
|||
|
|
@ -26,9 +26,10 @@ from typing import TYPE_CHECKING, Any
|
|||
|
||||
from agents import RunContextWrapper, function_tool
|
||||
|
||||
from strix.tools.mcp.client import dispatch_mcp_call
|
||||
from strix.tools.mcp.client import _errored_tool_output
|
||||
from strix.tools.mcp.naming import namespaced_tool_name
|
||||
from strix.tools.mcp.registry import MCP_REGISTRY_CONTEXT_KEY, McpRegistry
|
||||
from strix.tools.mcp.session import McpConnectionUnavailableError
|
||||
|
||||
|
||||
if TYPE_CHECKING:
|
||||
|
|
@ -49,6 +50,13 @@ def _unknown_connection(connection: str, registry: McpRegistry) -> str:
|
|||
return f"Unknown MCP connection {connection!r}. Available connections: {available}."
|
||||
|
||||
|
||||
def _unavailable_connection(connection: str) -> str:
|
||||
return (
|
||||
f"MCP connection {connection!r} is unavailable: its live session failed and "
|
||||
"could not be reconnected, so it is unavailable for the rest of this run."
|
||||
)
|
||||
|
||||
|
||||
def _format_tool(tool: MCPTool) -> str:
|
||||
schema = json.dumps(tool.inputSchema or {"type": "object"}, indent=2, ensure_ascii=False)
|
||||
description = (tool.description or "").strip() or "(no description)"
|
||||
|
|
@ -70,6 +78,7 @@ async def list_mcps(ctx: RunContextWrapper) -> dict[str, Any]:
|
|||
registry = _registry_from_ctx(ctx)
|
||||
if registry is None or not registry:
|
||||
return {"connections": []}
|
||||
dead_by_name = {status.name: status.dead for status in registry.statuses()}
|
||||
return {
|
||||
"connections": [
|
||||
{
|
||||
|
|
@ -77,6 +86,7 @@ async def list_mcps(ctx: RunContextWrapper) -> dict[str, Any]:
|
|||
"name": summary.name,
|
||||
"description": summary.purpose,
|
||||
"tool_count": summary.tool_count,
|
||||
"dead": dead_by_name.get(summary.name, False),
|
||||
}
|
||||
for summary in registry.summaries()
|
||||
]
|
||||
|
|
@ -102,7 +112,10 @@ async def describe_mcp(ctx: RunContextWrapper, connection: str) -> str:
|
|||
entry = registry.get(connection)
|
||||
if entry is None:
|
||||
return _unknown_connection(connection, registry)
|
||||
tools = await entry.server.list_tools()
|
||||
try:
|
||||
tools = await entry.session.list_tools()
|
||||
except McpConnectionUnavailableError:
|
||||
return _unavailable_connection(connection)
|
||||
if not tools:
|
||||
return f"MCP connection {connection!r} offers no tools."
|
||||
header = f"MCP connection {connection!r} offers {len(tools)} tool(s):"
|
||||
|
|
@ -155,7 +168,10 @@ async def call_mcp(
|
|||
return invalid_arguments
|
||||
if arguments is not None and not isinstance(arguments, dict):
|
||||
return invalid_arguments
|
||||
available = await entry.server.list_tools()
|
||||
try:
|
||||
available = await entry.session.list_tools()
|
||||
except McpConnectionUnavailableError:
|
||||
return _errored_tool_output(_unavailable_connection(connection))
|
||||
valid_names = {mcp_tool.name for mcp_tool in available}
|
||||
if tool not in valid_names:
|
||||
offered = ", ".join(sorted(valid_names)) or "(none)"
|
||||
|
|
@ -164,8 +180,7 @@ async def call_mcp(
|
|||
f"Tools this connection offers: {offered}. "
|
||||
"Call describe_mcp for their input schemas."
|
||||
)
|
||||
return await dispatch_mcp_call(
|
||||
entry.server,
|
||||
return await entry.session.dispatch(
|
||||
tool,
|
||||
arguments or {},
|
||||
label=namespaced_tool_name(connection, tool),
|
||||
|
|
|
|||
|
|
@ -31,6 +31,8 @@ from agents.mcp import (
|
|||
)
|
||||
from mcp.client.stdio import stdio_client
|
||||
|
||||
from strix.tools.mcp.session import McpConnectionUnavailableError, SupervisedMcpSession
|
||||
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from collections.abc import Callable
|
||||
|
|
@ -52,17 +54,18 @@ logger = logging.getLogger(__name__)
|
|||
|
||||
|
||||
class ConnectedMcpServer(NamedTuple):
|
||||
"""One successfully connected MCP server and how many tools it offers.
|
||||
"""One successfully connected MCP connection and how many tools it offers.
|
||||
|
||||
``server`` is kept so the caller can clean it up when the run ends, and so
|
||||
the caller can hand the live session to the run's
|
||||
``session`` is the :class:`~strix.tools.mcp.session.SupervisedMcpSession` that
|
||||
owns the live connection on its own task, so the caller cleans it up when the
|
||||
run ends (``await session.aclose()``) and hands it to the run's
|
||||
:class:`~strix.tools.mcp.registry.McpRegistry`; ``name`` and ``tool_count``
|
||||
let the caller show the user a startup summary and fill the prompt inventory;
|
||||
``notes`` carries the connection's optional free-text description so the
|
||||
caller can surface it as the connection's purpose in the inventory.
|
||||
"""
|
||||
|
||||
server: MCPServer
|
||||
session: SupervisedMcpSession
|
||||
name: str
|
||||
tool_count: int
|
||||
notes: str | None = None
|
||||
|
|
@ -227,66 +230,75 @@ def _errored_tool_output(tool_output: Any) -> dict[str, Any]:
|
|||
return {"success": False, "content": tool_output}
|
||||
|
||||
|
||||
async def _count_server_tools(config: McpConnectionConfig, server: MCPServer) -> int:
|
||||
"""Count a connected server's reachable tools for the startup summary.
|
||||
async def _count_session_tools(config: McpConnectionConfig, session: SupervisedMcpSession) -> int:
|
||||
"""Count a connected session's reachable tools for the startup summary.
|
||||
|
||||
``allowed_tools`` of ``None`` counts every listed tool; a list counts only
|
||||
those names. The count matches what ``describe_mcp`` will show, because the
|
||||
static tool filter built in :func:`_build_server` restricts the server's own
|
||||
``list_tools`` to the same allowlist.
|
||||
``list_tools`` to the same allowlist. The listing goes through the session's
|
||||
owning task like every other call.
|
||||
"""
|
||||
allowed = config.allowed_tools
|
||||
mcp_tools = await server.list_tools()
|
||||
mcp_tools = await session.list_tools()
|
||||
return sum(1 for mcp_tool in mcp_tools if allowed is None or mcp_tool.name in allowed)
|
||||
|
||||
|
||||
async def connect_mcp_servers(
|
||||
configs: list[McpConnectionConfig],
|
||||
) -> list[ConnectedMcpServer]:
|
||||
"""Connect to each MCP server and return its live session.
|
||||
"""Connect each MCP config on its own supervising task and return the sessions.
|
||||
|
||||
Returns one :class:`ConnectedMcpServer` per server that connected, carrying
|
||||
the SDK server (so the caller can clean it up when the run ends and hand it to
|
||||
the run's registry) plus the server name, how many tools it offers, and the
|
||||
connection's notes. Connections that fail are skipped rather than raised.
|
||||
Each connection becomes a :class:`~strix.tools.mcp.session.SupervisedMcpSession`
|
||||
that owns ``connect()``, the held-open session, and ``cleanup()`` on one
|
||||
dedicated task, so a later background failure in one session is contained to
|
||||
that task and never cancels the run. Returns one :class:`ConnectedMcpServer`
|
||||
per session that connected, carrying the session (the caller closes it with
|
||||
``await session.aclose()`` when the run ends and hands it to the run's
|
||||
registry) plus the connection name, tool count, and notes. A connection whose
|
||||
initial connect fails is skipped rather than raised (fail-open).
|
||||
|
||||
If this coroutine is itself cancelled mid-attach (the run going down), every
|
||||
session started so far is closed on its own task before the cancellation is
|
||||
re-raised, so nothing is orphaned.
|
||||
|
||||
Nothing is registered as an agent tool: the caller builds a per-run
|
||||
:class:`~strix.tools.mcp.registry.McpRegistry` from these sessions, and the
|
||||
agent reaches each tool on demand through ``describe_mcp`` / ``call_mcp``.
|
||||
"""
|
||||
connected: list[ConnectedMcpServer] = []
|
||||
for config in configs:
|
||||
server: MCPServer | None = None
|
||||
try:
|
||||
server = _build_server(config)
|
||||
await server.connect() # type: ignore[no-untyped-call]
|
||||
tool_count = await _count_server_tools(config, server)
|
||||
except Exception:
|
||||
logger.exception("Skipping MCP connection %r", config.name)
|
||||
if server is not None:
|
||||
with contextlib.suppress(Exception):
|
||||
await server.cleanup() # type: ignore[no-untyped-call]
|
||||
continue
|
||||
except BaseException:
|
||||
# A cancellation (or other non-Exception failure) mid-connect must not
|
||||
# orphan MCP subprocesses or HTTP sessions. Clean up the server being
|
||||
# connected and every server already connected, then re-raise so the
|
||||
# caller still stops. The runner only receives the list on a clean
|
||||
# return, so on an abnormal exit this function owns the cleanup.
|
||||
if server is not None:
|
||||
with contextlib.suppress(Exception):
|
||||
await server.cleanup() # type: ignore[no-untyped-call]
|
||||
for established in connected:
|
||||
with contextlib.suppress(Exception):
|
||||
await established.server.cleanup() # type: ignore[no-untyped-call]
|
||||
raise
|
||||
|
||||
logger.info("Connected MCP server %r (%d tools)", config.name, tool_count)
|
||||
connected.append(
|
||||
ConnectedMcpServer(
|
||||
server=server, name=config.name, tool_count=tool_count, notes=config.notes
|
||||
sessions: list[SupervisedMcpSession] = []
|
||||
try:
|
||||
for config in configs:
|
||||
session = SupervisedMcpSession(config)
|
||||
sessions.append(session)
|
||||
if not await session.start():
|
||||
# Initial connect failed; already logged inside the session. Drop it.
|
||||
await session.aclose()
|
||||
sessions.remove(session)
|
||||
continue
|
||||
try:
|
||||
tool_count = await _count_session_tools(config, session)
|
||||
except McpConnectionUnavailableError:
|
||||
# The session died between connecting and its first listing; skip it.
|
||||
logger.warning("MCP connection %r died before its first listing", config.name)
|
||||
await session.aclose()
|
||||
sessions.remove(session)
|
||||
continue
|
||||
logger.info("Connected MCP server %r (%d tools)", config.name, tool_count)
|
||||
connected.append(
|
||||
ConnectedMcpServer(
|
||||
session=session, name=config.name, tool_count=tool_count, notes=config.notes
|
||||
)
|
||||
)
|
||||
)
|
||||
except BaseException:
|
||||
# Cancelled or errored mid-attach: close every session started so far,
|
||||
# each on its own task, then re-raise. The runner only receives the list
|
||||
# on a clean return, so on an abnormal exit this function owns the cleanup.
|
||||
for session in sessions:
|
||||
with contextlib.suppress(BaseException):
|
||||
await session.aclose()
|
||||
raise
|
||||
|
||||
return connected
|
||||
|
||||
|
|
@ -318,7 +330,7 @@ async def attach_mcp_requests(
|
|||
request = request_by_name[connection.name]
|
||||
registry.add(
|
||||
name=connection.name,
|
||||
server=connection.server,
|
||||
session=connection.session,
|
||||
tool_count=connection.tool_count,
|
||||
purpose=request.purpose or connection.notes,
|
||||
provider=request.provider,
|
||||
|
|
|
|||
|
|
@ -25,6 +25,8 @@ from __future__ import annotations
|
|||
import dataclasses
|
||||
from typing import TYPE_CHECKING, Any, NamedTuple
|
||||
|
||||
from strix.tools.mcp.session import SupervisedMcpSession
|
||||
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from agents.mcp import MCPServer
|
||||
|
|
@ -54,24 +56,45 @@ MCP_DISPATCH_TOOLS = frozenset({CALL_MCP_TOOL, DESCRIBE_MCP_TOOL})
|
|||
class McpConnectionEntry:
|
||||
"""One live MCP connection a scan may reach, keyed by ``name``.
|
||||
|
||||
``server`` is the connected SDK session the dispatch tools list tools on and
|
||||
call tools through. ``purpose`` is the human label ``list_mcps`` reports as the
|
||||
connection's description (the user's connection notes, or whatever the caller
|
||||
supplies). ``tool_count`` is how many tools the connection offers, also
|
||||
reported by ``list_mcps``. ``result_transform``, when set, runs on each call's structured result
|
||||
at the single dispatch point (strix-pro's sanitizer uses it). ``provider`` is
|
||||
an optional source label (e.g. ``"supabase"``) the caller tags the connection
|
||||
with; the command-line path leaves it ``None``, and event tagging surfaces it
|
||||
when set.
|
||||
``session`` is the :class:`~strix.tools.mcp.session.SupervisedMcpSession` that
|
||||
owns the connection on its own task; the dispatch tools list tools and call
|
||||
tools through it (``session.list_tools`` / ``session.dispatch``) so a session
|
||||
failure is contained and can reconnect. ``purpose`` is the human label
|
||||
``list_mcps`` reports as the connection's description (the user's connection
|
||||
notes, or whatever the caller supplies). ``tool_count`` is how many tools the
|
||||
connection offers, also reported by ``list_mcps``. ``result_transform``, when
|
||||
set, runs on each call's structured result at the single dispatch point
|
||||
(strix-pro's sanitizer uses it). ``provider`` is an optional source label
|
||||
(e.g. ``"supabase"``) the caller tags the connection with; the command-line
|
||||
path leaves it ``None``, and event tagging surfaces it when set.
|
||||
|
||||
The connection config the session reconnects with (and its bearer token) lives
|
||||
on ``session`` in memory only. It is reached via :attr:`config` for the
|
||||
reconnect path and is never logged, serialized into the event stream, or
|
||||
written to disk.
|
||||
"""
|
||||
|
||||
server: MCPServer
|
||||
session: SupervisedMcpSession
|
||||
name: str
|
||||
purpose: str | None = None
|
||||
tool_count: int = 0
|
||||
result_transform: ResultTransform | None = None
|
||||
provider: str | None = None
|
||||
|
||||
@property
|
||||
def server(self) -> MCPServer | None:
|
||||
"""The current live server behind the session (swapped on reconnect).
|
||||
|
||||
Kept so existing callers that read ``entry.server`` keep working; new code
|
||||
should call through ``entry.session`` so reconnect and containment apply.
|
||||
"""
|
||||
return self.session.server
|
||||
|
||||
@property
|
||||
def config(self) -> McpConnectionConfig | None:
|
||||
"""The session's reconnect config. Carries the bearer token; never log it."""
|
||||
return self.session.config
|
||||
|
||||
|
||||
@dataclasses.dataclass(frozen=True)
|
||||
class McpConnectionSummary:
|
||||
|
|
@ -84,6 +107,24 @@ class McpConnectionSummary:
|
|||
provider: str | None = None
|
||||
|
||||
|
||||
@dataclasses.dataclass(frozen=True)
|
||||
class McpConnectionStatus:
|
||||
"""One connection's live status for the interfaces (the TUI panel, the app
|
||||
strip, and the roster signal the app consumes).
|
||||
|
||||
Non-secret by construction: only the connection ``name``, its ``provider``
|
||||
label, its ``tool_count``, and whether its live session is currently ``dead``
|
||||
(its reconnect-retry gave up). No config, token, url, or purpose rides here.
|
||||
``dead`` is read live off the connection's session at the moment this is
|
||||
built, so a fresh :meth:`McpRegistry.statuses` reflects the current health.
|
||||
"""
|
||||
|
||||
name: str
|
||||
provider: str | None
|
||||
tool_count: int
|
||||
dead: bool
|
||||
|
||||
|
||||
@dataclasses.dataclass(frozen=True)
|
||||
class McpConnectionRequest:
|
||||
"""A source-agnostic request to attach one MCP connection to a run.
|
||||
|
|
@ -129,15 +170,28 @@ class McpRegistry:
|
|||
self,
|
||||
*,
|
||||
name: str,
|
||||
server: MCPServer,
|
||||
session: SupervisedMcpSession | None = None,
|
||||
server: MCPServer | None = None,
|
||||
config: McpConnectionConfig | None = None,
|
||||
purpose: str | None = None,
|
||||
tool_count: int = 0,
|
||||
result_transform: ResultTransform | None = None,
|
||||
provider: str | None = None,
|
||||
) -> McpConnectionEntry:
|
||||
"""Register one connection under ``name`` (last write wins)."""
|
||||
"""Register one connection under ``name`` (last write wins).
|
||||
|
||||
Pass ``session`` for a session the engine already supervises (the attach
|
||||
path does this). Pass ``server`` for an already-connected server the caller
|
||||
owns (strix-pro's cloud sessions): it is adopted into a session that runs
|
||||
calls inline against it, and reconnects only when a ``config`` is also
|
||||
given. Exactly one of ``session`` or ``server`` is required.
|
||||
"""
|
||||
if session is None:
|
||||
if server is None:
|
||||
raise ValueError("McpRegistry.add requires either 'session' or 'server'")
|
||||
session = SupervisedMcpSession.adopt(server, name=name, config=config)
|
||||
entry = McpConnectionEntry(
|
||||
server=server,
|
||||
session=session,
|
||||
name=name,
|
||||
purpose=purpose,
|
||||
tool_count=tool_count,
|
||||
|
|
@ -167,6 +221,23 @@ class McpRegistry:
|
|||
for entry in self._entries.values()
|
||||
]
|
||||
|
||||
def statuses(self) -> list[McpConnectionStatus]:
|
||||
"""One live status per connection, in insertion order.
|
||||
|
||||
Reads each connection's ``dead`` flag off its session at call time, so the
|
||||
interfaces (the TUI panel via the Python backend projection, and the
|
||||
roster signal the app consumes) get the current health each time they
|
||||
rebuild. Non-secret: name, provider, tool_count, dead only."""
|
||||
return [
|
||||
McpConnectionStatus(
|
||||
name=entry.name,
|
||||
provider=entry.provider,
|
||||
tool_count=entry.tool_count,
|
||||
dead=entry.session.is_dead,
|
||||
)
|
||||
for entry in self._entries.values()
|
||||
]
|
||||
|
||||
def clear(self) -> None:
|
||||
"""Drop every connection (the sessions themselves are closed by the
|
||||
runner)."""
|
||||
|
|
|
|||
536
strix/tools/mcp/session.py
Normal file
536
strix/tools/mcp/session.py
Normal file
|
|
@ -0,0 +1,536 @@
|
|||
"""Own each MCP connection's live session on its own supervising task.
|
||||
|
||||
The bug this fixes: the streamable-HTTP transport (the ``mcp`` SDK) opens an
|
||||
internal anyio task group when ``server.connect()`` runs, and that task group's
|
||||
cancel scope is entered on whatever task called ``connect()`` and stays open for
|
||||
the session's whole life. In the old code that task was the run's main task, the
|
||||
one the agent loop runs on. So when a provider returned an HTTP error on one of
|
||||
the transport's background tasks (for example a ``403`` on a background POST),
|
||||
the task group cancelled its scope, the cancellation
|
||||
propagated to the main task, and the whole scan died with a bare
|
||||
``CancelledError`` (mislabeled as a user interrupt). Teardown then raised
|
||||
"Attempted to exit cancel scope in a different task than it was entered in"
|
||||
because cleanup ran on a different task than connect.
|
||||
|
||||
The fix, mirroring how child agents run on their own ``asyncio.create_task``
|
||||
(see :func:`strix.core.execution.spawn_child_agent`): give each connection its
|
||||
own dedicated supervising task that owns ``connect()``, the session's held-open
|
||||
lifetime, and ``cleanup()``. Three consequences:
|
||||
|
||||
- **Containment.** The transport's cancel scope is now entered on the supervising
|
||||
task, so a background failure cancels only that task. The run and every other
|
||||
connection keep going.
|
||||
- **Co-located teardown.** ``connect()`` and ``cleanup()`` run on the same task,
|
||||
so the "exit cancel scope in a different task" error cannot happen.
|
||||
- **A value, not a cancellation, reaches the caller.** The agent never touches the
|
||||
live session directly. It hands a call to the supervising task over a queue and
|
||||
awaits the result as a value; if the session task dies, the caller gets a
|
||||
"connection unavailable" value instead of a cancellation propagating into the
|
||||
agent loop.
|
||||
|
||||
When a call fails the supervisor rebuilds and reconnects the session once (reusing
|
||||
the same config, so the same bearer token, never re-fetching credentials) and
|
||||
re-runs the one failed call once. If that still fails, the connection is marked
|
||||
dead: every later call returns the standard failed-tool output.
|
||||
|
||||
Security: the connection's :class:`~strix.tools.mcp.config.McpConnectionConfig`
|
||||
holds a live bearer credential and is kept here in memory only, on the same
|
||||
in-process object that already holds the live session. It is never logged,
|
||||
serialized into the run's event stream, or written to disk; :meth:`__repr__`
|
||||
omits it and the token field's own ``repr`` is already suppressed.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import asyncio
|
||||
import contextlib
|
||||
import dataclasses
|
||||
import logging
|
||||
from typing import TYPE_CHECKING, Any, cast
|
||||
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from collections.abc import Awaitable, Callable
|
||||
|
||||
from agents.mcp import MCPServer
|
||||
from mcp.types import Tool as MCPTool
|
||||
|
||||
from strix.tools.mcp.client import ResultTransform
|
||||
from strix.tools.mcp.config import McpConnectionConfig
|
||||
|
||||
# One operation to run against the live session, e.g. ``list_tools`` or a tool
|
||||
# call. Runs on the supervising task (supervised sessions) or inline (adopted
|
||||
# sessions), and its return value becomes the caller's result.
|
||||
Job = Callable[[MCPServer], Awaitable[Any]]
|
||||
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
# How long a graceful (sentinel) shutdown waits for the serve loop to drain
|
||||
# before the supervising task is cancelled instead. Bounds teardown so a slow or
|
||||
# hung in-flight call cannot stall it forever.
|
||||
_SHUTDOWN_TIMEOUT = 10.0
|
||||
|
||||
|
||||
class McpConnectionUnavailableError(RuntimeError):
|
||||
"""A dead MCP connection could not be reached and did not come back.
|
||||
|
||||
Raised by :meth:`SupervisedMcpSession.list_tools` when the connection is dead
|
||||
so the read-only dispatch tools (``describe_mcp``) can report it cleanly.
|
||||
:meth:`SupervisedMcpSession.dispatch` does not raise it: a call to a dead
|
||||
connection returns the standard failed-tool output instead.
|
||||
"""
|
||||
|
||||
|
||||
@dataclasses.dataclass
|
||||
class _Outcome:
|
||||
"""What running one job resolved to: a value, or the connection being dead."""
|
||||
|
||||
value: Any = None
|
||||
dead: bool = False
|
||||
|
||||
|
||||
@dataclasses.dataclass
|
||||
class _Request:
|
||||
"""One job handed to the supervising task, with the future its result lands in."""
|
||||
|
||||
job: Job
|
||||
future: asyncio.Future[_Outcome]
|
||||
|
||||
|
||||
class SupervisedMcpSession:
|
||||
"""One MCP connection whose live session is owned by a dedicated task.
|
||||
|
||||
Built two ways:
|
||||
|
||||
- :meth:`__init__` + :meth:`start` for a *supervised* session: the engine owns
|
||||
connecting. ``start`` spawns the supervising task, which builds and connects
|
||||
the server on itself and then serves calls handed to it over a queue. This is
|
||||
the path that contains a background session failure to one task.
|
||||
- :meth:`adopt` for an *adopted* session: the caller already holds a connected
|
||||
server (strix-pro's cloud sessions, and the test fakes). There is no
|
||||
supervising task; calls run inline against the given server. Reconnect works
|
||||
only when a config was supplied.
|
||||
|
||||
Public async API used by the dispatch tools: :meth:`list_tools` and
|
||||
:meth:`dispatch`. Lifecycle: :meth:`start`, :meth:`aclose`. Read-only:
|
||||
:attr:`name`, :attr:`server`, :attr:`config`, :attr:`is_dead`.
|
||||
"""
|
||||
|
||||
def __init__(self, config: McpConnectionConfig) -> None:
|
||||
self._name = config.name
|
||||
self._config: McpConnectionConfig | None = config
|
||||
self._server: MCPServer | None = None
|
||||
self._supervised = True
|
||||
self._task: asyncio.Task[None] | None = None
|
||||
self._queue: asyncio.Queue[_Request | None] | None = None
|
||||
self._ready: asyncio.Future[bool] | None = None
|
||||
self._pending: set[asyncio.Future[_Outcome]] = set()
|
||||
self._dead = False
|
||||
self._closing = False
|
||||
self._on_dead: Callable[[], None] | None = None
|
||||
# Guards the idle-death self-heal against a flapping server: set after an
|
||||
# idle reconnect, cleared once a real call runs. If the session dies idle
|
||||
# again before serving anything, we give up instead of reconnecting in a
|
||||
# tight loop.
|
||||
self._healed_without_progress = False
|
||||
|
||||
@classmethod
|
||||
def adopt(
|
||||
cls,
|
||||
server: MCPServer,
|
||||
*,
|
||||
name: str,
|
||||
config: McpConnectionConfig | None = None,
|
||||
) -> SupervisedMcpSession:
|
||||
"""Wrap an already-connected server without a supervising task.
|
||||
|
||||
Calls run inline against ``server`` on the caller's task, matching the old
|
||||
direct-dispatch behavior. Reconnect is available only when ``config`` is
|
||||
given; otherwise a failed call marks the connection dead.
|
||||
"""
|
||||
self = cls.__new__(cls)
|
||||
self._name = name
|
||||
self._config = config
|
||||
self._server = server
|
||||
self._supervised = False
|
||||
self._task = None
|
||||
self._queue = None
|
||||
self._ready = None
|
||||
self._pending = set()
|
||||
self._dead = False
|
||||
self._closing = False
|
||||
self._on_dead = None
|
||||
self._healed_without_progress = False
|
||||
return self
|
||||
|
||||
# -- read-only accessors --------------------------------------------------
|
||||
|
||||
@property
|
||||
def name(self) -> str:
|
||||
return self._name
|
||||
|
||||
@property
|
||||
def server(self) -> MCPServer | None:
|
||||
"""The current live server, or ``None`` once dead. Swapped on reconnect."""
|
||||
return self._server
|
||||
|
||||
@property
|
||||
def config(self) -> McpConnectionConfig | None:
|
||||
"""The connection config kept for reconnect. Carries the bearer token, so
|
||||
never log or serialize this."""
|
||||
return self._config
|
||||
|
||||
@property
|
||||
def is_dead(self) -> bool:
|
||||
return self._dead
|
||||
|
||||
def set_on_dead(self, callback: Callable[[], None] | None) -> None:
|
||||
"""Register a one-shot callback fired when the connection transitions to dead.
|
||||
|
||||
The callback runs on whatever task marks the connection dead (the
|
||||
supervising task for a supervised session, the caller's task for an
|
||||
adopted one), so it must not block. It fires at most once, on the
|
||||
healthy->dead edge, and never for a connection that only ever shut down
|
||||
cleanly. The interfaces use it to push a live "offline" status without
|
||||
polling. Exceptions from the callback are swallowed (logged) so a status
|
||||
push can never take down the session task.
|
||||
"""
|
||||
self._on_dead = callback
|
||||
|
||||
def _mark_dead(self) -> None:
|
||||
"""Flip the connection to dead and fire ``on_dead`` once on the transition."""
|
||||
if self._dead:
|
||||
return
|
||||
self._dead = True
|
||||
callback = self._on_dead
|
||||
if callback is None:
|
||||
return
|
||||
try:
|
||||
callback()
|
||||
except Exception:
|
||||
logger.exception("MCP on_dead callback for %r failed", self._name)
|
||||
|
||||
def __repr__(self) -> str:
|
||||
# Deliberately omits the config so the bearer token can never reach a log
|
||||
# line through an accidental repr of this object.
|
||||
return f"SupervisedMcpSession(name={self._name!r}, dead={self._dead})"
|
||||
|
||||
# -- lifecycle ------------------------------------------------------------
|
||||
|
||||
async def start(self) -> bool:
|
||||
"""Spawn the supervising task, connect on it, and wait until it is ready.
|
||||
|
||||
Returns ``True`` when the session connected, ``False`` when the initial
|
||||
connect failed (the caller then skips this connection, fail-open). Only
|
||||
valid for a supervised session.
|
||||
"""
|
||||
loop = asyncio.get_running_loop()
|
||||
self._queue = asyncio.Queue()
|
||||
self._ready = loop.create_future()
|
||||
self._task = asyncio.create_task(self._supervise(), name=f"mcp-session-{self._name}")
|
||||
return await self._ready
|
||||
|
||||
async def aclose(self) -> None:
|
||||
"""Shut the connection down and clean up its session on its owning task.
|
||||
|
||||
For a connected supervised session this signals the supervising task with a
|
||||
sentinel so ``cleanup()`` runs on the same task that ran ``connect()``,
|
||||
giving an orderly shutdown the supervisor tells apart from a session death.
|
||||
Teardown is always bounded: if the serve loop cannot drain the sentinel in
|
||||
time (a slow or hung in-flight call), or the session never finished
|
||||
connecting (including a connect cancelled mid-await), the task is cancelled
|
||||
instead. ``_closing`` is set first, so the supervisor treats that
|
||||
cancellation as shutdown and still cleans up on its own task.
|
||||
"""
|
||||
self._closing = True
|
||||
if self._supervised and self._task is not None:
|
||||
if not self._task.done():
|
||||
# A cancelled readiness future (the connect was cancelled mid-await)
|
||||
# counts as "not connected": never call ``.result()`` on it, which
|
||||
# would raise here and skip the cleanup below.
|
||||
connected = (
|
||||
self._ready is not None
|
||||
and self._ready.done()
|
||||
and not self._ready.cancelled()
|
||||
and self._ready.result()
|
||||
)
|
||||
if connected and self._queue is not None:
|
||||
# Reached the serve loop: a sentinel gives a clean, cancel-free
|
||||
# teardown, with cleanup() running on the supervising task. Bound
|
||||
# it, though: a hung in-flight call would otherwise leave the
|
||||
# sentinel queued behind it forever, so cancel the task if the
|
||||
# drain does not finish in time (wait_for cancels it on timeout).
|
||||
with contextlib.suppress(Exception):
|
||||
await self._queue.put(None)
|
||||
with contextlib.suppress(
|
||||
asyncio.TimeoutError, asyncio.CancelledError, Exception
|
||||
):
|
||||
await asyncio.wait_for(self._task, _SHUTDOWN_TIMEOUT)
|
||||
else:
|
||||
# Still stuck in connect(), never connected, or connect
|
||||
# cancelled: cancel to unstick it.
|
||||
self._task.cancel()
|
||||
with contextlib.suppress(asyncio.CancelledError, Exception):
|
||||
await self._task
|
||||
else:
|
||||
await self._safe_cleanup()
|
||||
self._fail_pending()
|
||||
|
||||
# -- caller-facing operations --------------------------------------------
|
||||
|
||||
async def list_tools(self) -> list[MCPTool]:
|
||||
"""List the connection's tools, reconnecting once if the session died.
|
||||
|
||||
Raises :class:`McpConnectionUnavailableError` when the connection is dead.
|
||||
"""
|
||||
outcome = await self._run_job(lambda server: server.list_tools())
|
||||
if outcome.dead:
|
||||
raise McpConnectionUnavailableError(self._unavailable_message())
|
||||
return cast("list[MCPTool]", outcome.value)
|
||||
|
||||
async def dispatch(
|
||||
self,
|
||||
tool_name: str,
|
||||
arguments: dict[str, Any],
|
||||
*,
|
||||
label: str,
|
||||
result_transform: ResultTransform | None = None,
|
||||
) -> Any:
|
||||
"""Run one tool call, reconnecting once and retrying once on session death.
|
||||
|
||||
Returns the tool output on success, or the standard failed-tool output
|
||||
(``success: False``) with a "connection unavailable" message when the
|
||||
connection is dead.
|
||||
"""
|
||||
from strix.tools.mcp.client import dispatch_mcp_call
|
||||
|
||||
async def job(server: MCPServer) -> Any:
|
||||
return await dispatch_mcp_call(
|
||||
server,
|
||||
tool_name,
|
||||
arguments,
|
||||
label=label,
|
||||
result_transform=result_transform,
|
||||
)
|
||||
|
||||
outcome = await self._run_job(job)
|
||||
if outcome.dead:
|
||||
from strix.tools.mcp.client import _errored_tool_output
|
||||
|
||||
return _errored_tool_output(self._unavailable_message())
|
||||
return outcome.value
|
||||
|
||||
# -- job routing ----------------------------------------------------------
|
||||
|
||||
async def _run_job(self, job: Job) -> _Outcome:
|
||||
"""Route one job to the owning task (supervised) or run it inline (adopted)."""
|
||||
if self._supervised:
|
||||
return await self._submit(job)
|
||||
return await self._execute(job)
|
||||
|
||||
async def _submit(self, job: Job) -> _Outcome:
|
||||
"""Hand a job to the supervising task and await its result as a value."""
|
||||
if self._dead or self._closing or self._task is None or self._task.done():
|
||||
return _Outcome(dead=True)
|
||||
loop = asyncio.get_running_loop()
|
||||
future: asyncio.Future[_Outcome] = loop.create_future()
|
||||
self._pending.add(future)
|
||||
if self._queue is None:
|
||||
self._pending.discard(future)
|
||||
return _Outcome(dead=True)
|
||||
await self._queue.put(_Request(job=job, future=future))
|
||||
# The task may have ended between the guard above and the put; ``_fail_pending``
|
||||
# would then never see this future, so resolve it here.
|
||||
if self._task.done() and not future.done():
|
||||
self._pending.discard(future)
|
||||
return _Outcome(dead=True)
|
||||
return await future
|
||||
|
||||
# -- the supervising task -------------------------------------------------
|
||||
|
||||
async def _supervise(self) -> None:
|
||||
"""Own the session for its whole life on one task: connect, serve, clean up."""
|
||||
try:
|
||||
self._server = await self._open()
|
||||
except asyncio.CancelledError:
|
||||
# The connect was cancelled (the run is going down, or the transport
|
||||
# scope cancelled mid-connect). Report not-ready so the attach path
|
||||
# treats it as a skipped connection; do not propagate.
|
||||
self._report_ready(value=False)
|
||||
await self._safe_cleanup()
|
||||
self._fail_pending()
|
||||
return
|
||||
except Exception:
|
||||
logger.exception("Skipping MCP connection %r", self._name)
|
||||
self._report_ready(value=False)
|
||||
await self._safe_cleanup()
|
||||
self._fail_pending()
|
||||
return
|
||||
|
||||
self._report_ready(value=True)
|
||||
try:
|
||||
await self._serve_loop()
|
||||
finally:
|
||||
await self._safe_cleanup()
|
||||
self._fail_pending()
|
||||
|
||||
async def _serve_loop(self) -> None:
|
||||
assert self._queue is not None
|
||||
while True:
|
||||
try:
|
||||
request = await self._queue.get()
|
||||
except asyncio.CancelledError:
|
||||
# A cancellation while idle is the transport's task group cancelling
|
||||
# this supervising task because a background session task failed.
|
||||
# Contained here. If we are closing, this is an ordinary shutdown,
|
||||
# so let it propagate. Otherwise try to self-heal once: reconnect a
|
||||
# fresh session and keep serving. The flag stops a flapping server
|
||||
# (one that dies again before serving any call) from reconnecting in
|
||||
# a tight loop; there we give up and mark the connection dead. Later
|
||||
# calls then short-circuit to the dead output without this task.
|
||||
if self._closing:
|
||||
raise
|
||||
if not self._healed_without_progress:
|
||||
logger.warning(
|
||||
"MCP connection %r session died while idle; reconnecting once",
|
||||
self._name,
|
||||
)
|
||||
if await self._reconnect():
|
||||
logger.info(
|
||||
"MCP connection %r reconnected after an idle death", self._name
|
||||
)
|
||||
self._healed_without_progress = True
|
||||
continue
|
||||
else:
|
||||
logger.warning(
|
||||
"MCP connection %r died again before serving a call; "
|
||||
"marking it unavailable",
|
||||
self._name,
|
||||
)
|
||||
self._mark_dead()
|
||||
await self._safe_cleanup()
|
||||
return
|
||||
if request is None: # shutdown sentinel
|
||||
return
|
||||
outcome = await self._execute(request.job)
|
||||
# A served call is real progress: clear the idle-heal guard so a future
|
||||
# idle death is again allowed one reconnect.
|
||||
self._healed_without_progress = False
|
||||
if not request.future.done():
|
||||
request.future.set_result(outcome)
|
||||
self._pending.discard(request.future)
|
||||
|
||||
# -- run one job with reconnect-once + retry-once -------------------------
|
||||
|
||||
async def _execute(self, job: Job) -> _Outcome:
|
||||
"""Run one job; on a session failure reconnect once and retry it once."""
|
||||
if self._dead or self._server is None:
|
||||
return _Outcome(dead=True)
|
||||
try:
|
||||
return _Outcome(value=await job(self._server))
|
||||
except asyncio.CancelledError:
|
||||
# For a supervised session a cancellation here is the transport scope
|
||||
# dying under an in-flight call: a session death, not a real cancel
|
||||
# (shutdown never cancels the task, it uses the sentinel). For an
|
||||
# adopted session there is no such scope, so a cancel is real.
|
||||
if not self._supervised or self._closing:
|
||||
raise
|
||||
logger.warning(
|
||||
"MCP connection %r was cancelled mid-call (session died); reconnecting once",
|
||||
self._name,
|
||||
)
|
||||
except Exception: # noqa: BLE001 - any call failure is treated as a session death
|
||||
logger.warning(
|
||||
"MCP connection %r failed mid-call; reconnecting once", self._name
|
||||
)
|
||||
|
||||
if not await self._reconnect():
|
||||
self._mark_dead()
|
||||
return _Outcome(dead=True)
|
||||
|
||||
try:
|
||||
return _Outcome(value=await job(self._server))
|
||||
except asyncio.CancelledError:
|
||||
if not self._supervised or self._closing:
|
||||
raise
|
||||
logger.warning(
|
||||
"MCP connection %r was cancelled again after reconnect; marking it unavailable",
|
||||
self._name,
|
||||
)
|
||||
self._mark_dead()
|
||||
return _Outcome(dead=True)
|
||||
except Exception: # noqa: BLE001 - any retry failure means the connection is dead
|
||||
logger.warning(
|
||||
"MCP connection %r failed again after reconnect; marking it unavailable",
|
||||
self._name,
|
||||
)
|
||||
self._mark_dead()
|
||||
return _Outcome(dead=True)
|
||||
|
||||
async def _reconnect(self) -> bool:
|
||||
"""Rebuild and reconnect the session once, reusing the stored config/token."""
|
||||
await self._safe_cleanup()
|
||||
if self._config is None:
|
||||
return False
|
||||
try:
|
||||
self._server = await self._open()
|
||||
except asyncio.CancelledError:
|
||||
if self._closing:
|
||||
raise
|
||||
logger.warning("MCP reconnect for %r was cancelled; giving up", self._name)
|
||||
self._server = None
|
||||
return False
|
||||
except Exception:
|
||||
logger.exception("MCP reconnect for %r failed", self._name)
|
||||
self._server = None
|
||||
return False
|
||||
logger.info("MCP connection %r reconnected", self._name)
|
||||
return True
|
||||
|
||||
async def _open(self) -> MCPServer:
|
||||
"""Build and connect the SDK server, reusing the existing setup steps.
|
||||
|
||||
If ``connect()`` fails, the just-built server is cleaned up here on this
|
||||
same task before the error propagates, so a failed connect never orphans
|
||||
an MCP subprocess or half-open HTTP session.
|
||||
"""
|
||||
from strix.tools.mcp.client import _build_server
|
||||
|
||||
if self._config is None:
|
||||
raise RuntimeError(f"MCP connection {self._name!r} has no config to connect")
|
||||
server = _build_server(self._config)
|
||||
try:
|
||||
await server.connect() # type: ignore[no-untyped-call]
|
||||
except BaseException:
|
||||
with contextlib.suppress(Exception):
|
||||
await server.cleanup() # type: ignore[no-untyped-call]
|
||||
raise
|
||||
return server
|
||||
|
||||
# -- helpers --------------------------------------------------------------
|
||||
|
||||
async def _safe_cleanup(self) -> None:
|
||||
server = self._server
|
||||
self._server = None
|
||||
if server is None:
|
||||
return
|
||||
with contextlib.suppress(Exception):
|
||||
await server.cleanup() # type: ignore[no-untyped-call]
|
||||
|
||||
def _report_ready(self, value: bool) -> None:
|
||||
if self._ready is not None and not self._ready.done():
|
||||
self._ready.set_result(value)
|
||||
|
||||
def _fail_pending(self) -> None:
|
||||
for future in self._pending:
|
||||
if not future.done():
|
||||
future.set_result(_Outcome(dead=True))
|
||||
self._pending.clear()
|
||||
|
||||
def _unavailable_message(self) -> str:
|
||||
return (
|
||||
f"MCP connection {self._name!r} is unavailable: its live session could "
|
||||
"not be reached and a reconnect attempt failed. It is marked unavailable "
|
||||
"for the rest of this run."
|
||||
)
|
||||
|
|
@ -8,6 +8,7 @@ two dispatch tools ``describe_mcp`` and ``call_mcp``.
|
|||
from __future__ import annotations
|
||||
|
||||
import asyncio
|
||||
import contextlib
|
||||
import json
|
||||
import re
|
||||
from typing import TYPE_CHECKING, Any
|
||||
|
|
@ -29,6 +30,7 @@ from strix.tools.mcp import (
|
|||
McpConnectionConfig,
|
||||
McpConnectionRequest,
|
||||
McpRegistry,
|
||||
SupervisedMcpSession,
|
||||
attach_mcp_requests,
|
||||
call_mcp,
|
||||
describe_mcp,
|
||||
|
|
@ -38,9 +40,11 @@ from strix.tools.mcp import (
|
|||
resolve_mcp_call,
|
||||
)
|
||||
from strix.tools.mcp import client as mcp_client
|
||||
from strix.tools.mcp import session as mcp_session_mod
|
||||
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from collections.abc import Callable
|
||||
from pathlib import Path
|
||||
|
||||
|
||||
|
|
@ -171,6 +175,13 @@ def _ctx(registry: McpRegistry | None) -> ToolContext[dict[str, Any]]:
|
|||
)
|
||||
|
||||
|
||||
async def _aclose_all(connections: list[Any]) -> None:
|
||||
"""Close every supervised session a connect/attach test opened, so no
|
||||
supervising task leaks into the event loop's teardown."""
|
||||
for connection in connections:
|
||||
await connection.session.aclose()
|
||||
|
||||
|
||||
@pytest.fixture(autouse=True)
|
||||
def _clear_mcp_env(monkeypatch: pytest.MonkeyPatch) -> None:
|
||||
"""Hide any MCP settings the developer has exported in their own shell."""
|
||||
|
|
@ -298,6 +309,8 @@ async def test_connect_returns_sessions_without_registering_agent_tools(
|
|||
assert [(c.name, c.tool_count) for c in connections] == [("fs", 2), ("db", 1)]
|
||||
assert list(factory.registered_agent_tools()) == before
|
||||
|
||||
await _aclose_all(connections)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_tool_count_honors_the_allowlist(monkeypatch: pytest.MonkeyPatch) -> None:
|
||||
|
|
@ -308,6 +321,8 @@ async def test_tool_count_honors_the_allowlist(monkeypatch: pytest.MonkeyPatch)
|
|||
|
||||
assert connections[0].tool_count == 1
|
||||
|
||||
await _aclose_all(connections)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_connection_notes_ride_on_the_connection(
|
||||
|
|
@ -326,6 +341,8 @@ async def test_connection_notes_ride_on_the_connection(
|
|||
|
||||
assert connections[0].notes == "Staging analytics DB; read-only."
|
||||
|
||||
await _aclose_all(connections)
|
||||
|
||||
|
||||
# --- server build branch -----------------------------------------------------
|
||||
|
||||
|
|
@ -399,11 +416,18 @@ async def test_list_mcps_returns_connections_with_ids_and_descriptions() -> None
|
|||
out = await list_mcps.on_invoke_tool(_ctx(registry), "{}")
|
||||
|
||||
# ``id`` is the exact connection name describe_mcp/call_mcp accept;
|
||||
# ``description`` is the summary's purpose; no tool schemas are included.
|
||||
# ``description`` is the summary's purpose; ``dead`` is the connection's live
|
||||
# health (both healthy here); no tool schemas are included.
|
||||
assert out == {
|
||||
"connections": [
|
||||
{"id": "fs", "name": "fs", "description": "local files", "tool_count": 2},
|
||||
{"id": "db", "name": "db", "description": None, "tool_count": 1},
|
||||
{
|
||||
"id": "fs",
|
||||
"name": "fs",
|
||||
"description": "local files",
|
||||
"tool_count": 2,
|
||||
"dead": False,
|
||||
},
|
||||
{"id": "db", "name": "db", "description": None, "tool_count": 1, "dead": False},
|
||||
]
|
||||
}
|
||||
|
||||
|
|
@ -801,37 +825,82 @@ def test_loader_exclude_selection_drops_named(
|
|||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_connect_cleans_up_when_cancelled_mid_connect(
|
||||
async def test_connect_skips_a_connection_whose_connect_is_cancelled(
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
) -> None:
|
||||
# Each connection now connects on its own supervising task. A cancellation of
|
||||
# one session's connect (the transport scope dying mid-connect) is contained
|
||||
# to that task: the connection is skipped and cleaned up, and the run's attach
|
||||
# keeps going rather than being cancelled.
|
||||
cleaned: list[str] = []
|
||||
|
||||
class _Tracking(FakeMCPServer):
|
||||
def __init__(self, name: str, *, fail_connect: bool = False) -> None:
|
||||
def __init__(self, name: str, *, cancel_connect: bool = False) -> None:
|
||||
super().__init__(name, [_mcp_tool("t")])
|
||||
self._fail_connect = fail_connect
|
||||
self._cancel_connect = cancel_connect
|
||||
|
||||
async def connect(self) -> None:
|
||||
if self._fail_connect:
|
||||
if self._cancel_connect:
|
||||
raise asyncio.CancelledError
|
||||
|
||||
async def cleanup(self) -> None:
|
||||
cleaned.append(self._name)
|
||||
|
||||
servers = {"good": _Tracking("good"), "bad": _Tracking("bad", fail_connect=True)}
|
||||
servers = {"good": _Tracking("good"), "bad": _Tracking("bad", cancel_connect=True)}
|
||||
monkeypatch.setattr(mcp_client, "_build_server", lambda config: servers[config.name])
|
||||
|
||||
configs = [
|
||||
McpConnectionConfig(name="good", url="https://mcp.example.com", allowed_tools=["t"]),
|
||||
McpConnectionConfig(name="bad", url="https://mcp.example.com", allowed_tools=["t"]),
|
||||
]
|
||||
configs = [_config("good", ["t"]), _config("bad", ["t"])]
|
||||
|
||||
connections = await mcp_client.connect_mcp_servers(configs)
|
||||
|
||||
# The cancelled connect is skipped and cleaned up; the good one is returned.
|
||||
assert [c.name for c in connections] == ["good"]
|
||||
assert "bad" in cleaned
|
||||
|
||||
await _aclose_all(connections)
|
||||
assert "good" in cleaned
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_connect_cleans_up_started_sessions_when_attach_is_cancelled(
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
) -> None:
|
||||
# If the attach coroutine itself is cancelled (the run going down) while a
|
||||
# later connection is still connecting, every session started so far is closed
|
||||
# on its own task before the cancellation is re-raised, so nothing is orphaned.
|
||||
cleaned: list[str] = []
|
||||
|
||||
class _Tracking(FakeMCPServer):
|
||||
def __init__(self, name: str, *, block_connect: bool = False) -> None:
|
||||
super().__init__(name, [_mcp_tool("t")])
|
||||
self._block_connect = block_connect
|
||||
|
||||
async def connect(self) -> None:
|
||||
if self._block_connect:
|
||||
await asyncio.Event().wait() # never completes
|
||||
|
||||
async def cleanup(self) -> None:
|
||||
cleaned.append(self._name)
|
||||
|
||||
servers = {"good": _Tracking("good"), "slow": _Tracking("slow", block_connect=True)}
|
||||
monkeypatch.setattr(mcp_client, "_build_server", lambda config: servers[config.name])
|
||||
|
||||
async def _attach() -> list[Any]:
|
||||
# Connect "good" first, then hang forever connecting "slow".
|
||||
return await mcp_client.connect_mcp_servers(
|
||||
[_config("good", ["t"]), _config("slow", ["t"])]
|
||||
)
|
||||
|
||||
task = asyncio.create_task(_attach())
|
||||
# Give the loop time to connect good and reach slow's hanging connect.
|
||||
for _ in range(100):
|
||||
await asyncio.sleep(0)
|
||||
task.cancel()
|
||||
with pytest.raises(asyncio.CancelledError):
|
||||
await mcp_client.connect_mcp_servers(configs)
|
||||
await task
|
||||
|
||||
# The server being connected when cancelled, and the one already connected,
|
||||
# are both cleaned up rather than orphaned.
|
||||
assert cleaned == ["bad", "good"]
|
||||
# The already-connected "good" session was cleaned up, not orphaned.
|
||||
assert "good" in cleaned
|
||||
|
||||
|
||||
# --- reading a tool call back to the server it went out to -------------------
|
||||
|
|
@ -916,6 +985,8 @@ async def test_attach_populates_registry_with_provider_and_transform(
|
|||
assert entry.purpose == "Customer DB"
|
||||
assert entry.result_transform is transform
|
||||
|
||||
await _aclose_all(connections)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_attach_bare_request_matches_the_command_line_shape(
|
||||
|
|
@ -933,7 +1004,7 @@ async def test_attach_bare_request_matches_the_command_line_shape(
|
|||
)
|
||||
|
||||
registry = McpRegistry()
|
||||
await attach_mcp_requests([McpConnectionRequest(config=config)], registry)
|
||||
connections = await attach_mcp_requests([McpConnectionRequest(config=config)], registry)
|
||||
|
||||
entry = registry.get("db")
|
||||
assert entry is not None
|
||||
|
|
@ -941,6 +1012,8 @@ async def test_attach_bare_request_matches_the_command_line_shape(
|
|||
assert entry.result_transform is None
|
||||
assert entry.purpose == "Staging analytics DB; read-only."
|
||||
|
||||
await _aclose_all(connections)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_attach_is_fail_open_and_skips_a_failed_connection(
|
||||
|
|
@ -970,6 +1043,8 @@ async def test_attach_is_fail_open_and_skips_a_failed_connection(
|
|||
assert registry.get("good") is not None
|
||||
assert registry.get("bad") is None
|
||||
|
||||
await _aclose_all(connections)
|
||||
|
||||
|
||||
# --- provider on the registry ------------------------------------------------
|
||||
|
||||
|
|
@ -1092,3 +1167,435 @@ async def test_errored_structured_output_is_wrapped_with_success_false() -> None
|
|||
# Structured content serializes to a JSON string; it too is wrapped under
|
||||
# ``content`` so the failure flag has a top-level dict to ride on.
|
||||
assert out == {"success": False, "content": json.dumps({"error": "boom"})}
|
||||
|
||||
|
||||
# --- per-session isolation: containment, reconnect-retry, mark-dead ----------
|
||||
# Each MCP connection's live session is owned by its own supervising task. A
|
||||
# background failure in one session is contained to that task: the agent's call
|
||||
# comes back as a value, the run keeps going, and the session reconnects once and
|
||||
# retries the failed call once before it is marked unavailable.
|
||||
|
||||
|
||||
class _DyingHttpServer(FakeMCPServer):
|
||||
"""A connected server whose ``call_tool`` fails to model a session death.
|
||||
|
||||
``death`` is the exception raised on a call: a plain ``Exception`` models an
|
||||
HTTP/transport error, and ``asyncio.CancelledError`` models the streamable-HTTP
|
||||
transport's task group cancelling the supervising task from a background POST
|
||||
error (for example a provider 403). ``alive`` flips to stop dying, so a
|
||||
reconnected replacement can succeed.
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
name: str,
|
||||
tools: list[MCPTool],
|
||||
*,
|
||||
death: BaseException,
|
||||
alive: bool = False,
|
||||
) -> None:
|
||||
super().__init__(name, tools)
|
||||
self._death = death
|
||||
self.alive = alive
|
||||
|
||||
async def call_tool(
|
||||
self,
|
||||
tool_name: str,
|
||||
arguments: dict[str, Any] | None,
|
||||
meta: dict[str, Any] | None = None,
|
||||
) -> CallToolResult:
|
||||
if not self.alive:
|
||||
raise self._death
|
||||
return await super().call_tool(tool_name, arguments)
|
||||
|
||||
|
||||
def _secret_config(name: str) -> McpConnectionConfig:
|
||||
return McpConnectionConfig(
|
||||
name=name,
|
||||
url="https://mcp.example.com",
|
||||
auth=BearerAuth(token="super-secret-bearer-token-42"),
|
||||
allowed_tools=["read_file"],
|
||||
)
|
||||
|
||||
|
||||
async def _started_session(config: McpConnectionConfig) -> SupervisedMcpSession:
|
||||
session = SupervisedMcpSession(config)
|
||||
assert await session.start()
|
||||
return session
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_call_mcp_reconnects_and_retries_after_a_session_death(
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
) -> None:
|
||||
# The first session dies on its call; the supervisor rebuilds the connection
|
||||
# once (reusing the existing _build_server + connect), retries the one call
|
||||
# once, and the retry lands on the healthy replacement.
|
||||
first = _DyingHttpServer("fs", [_mcp_tool("read_file")], death=ConnectionError("403"))
|
||||
second = FakeMCPServer("fs", [_mcp_tool("read_file")])
|
||||
built = iter([first, second])
|
||||
monkeypatch.setattr(mcp_client, "_build_server", lambda _config: next(built))
|
||||
|
||||
session = await _started_session(_secret_config("fs"))
|
||||
registry = McpRegistry()
|
||||
registry.add(name="fs", session=session, tool_count=1)
|
||||
|
||||
out = await call_mcp.on_invoke_tool(
|
||||
_ctx(registry), json.dumps({"connection": "fs", "tool": "read_file"})
|
||||
)
|
||||
|
||||
# The caller gets the tool output as a value, and the retried call ran on the
|
||||
# reconnected server.
|
||||
assert out == {"type": "text", "text": "routed:read_file"}
|
||||
assert second.calls == [("read_file", {})]
|
||||
assert session.is_dead is False
|
||||
|
||||
await session.aclose()
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_call_mcp_marks_connection_dead_when_reconnect_keeps_failing(
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
) -> None:
|
||||
# The session dies and the reconnect attempt also fails: the connection is
|
||||
# marked dead and the call returns the standard failed-tool output.
|
||||
first = _DyingHttpServer("fs", [_mcp_tool("read_file")], death=ConnectionError("403"))
|
||||
built = {"n": 0}
|
||||
|
||||
def _build(_config: McpConnectionConfig) -> MCPServer:
|
||||
built["n"] += 1
|
||||
if built["n"] == 1:
|
||||
return first
|
||||
raise ConnectionError("cannot reconnect")
|
||||
|
||||
monkeypatch.setattr(mcp_client, "_build_server", _build)
|
||||
|
||||
session = await _started_session(_secret_config("fs"))
|
||||
registry = McpRegistry()
|
||||
registry.add(name="fs", session=session, tool_count=1)
|
||||
|
||||
out = await call_mcp.on_invoke_tool(
|
||||
_ctx(registry), json.dumps({"connection": "fs", "tool": "read_file"})
|
||||
)
|
||||
|
||||
# A dead connection surfaces as an ordinary failed tool call, not an exception.
|
||||
assert isinstance(out, dict)
|
||||
assert out["success"] is False
|
||||
assert "unavailable" in out["content"]
|
||||
assert session.is_dead is True
|
||||
|
||||
# A later call short-circuits to the same failed output without a new attempt.
|
||||
again = await call_mcp.on_invoke_tool(
|
||||
_ctx(registry), json.dumps({"connection": "fs", "tool": "read_file"})
|
||||
)
|
||||
assert again["success"] is False
|
||||
|
||||
# describe_mcp reports the connection unavailable, and list_mcps still lists it.
|
||||
described = await describe_mcp.on_invoke_tool(_ctx(registry), json.dumps({"connection": "fs"}))
|
||||
assert "unavailable" in described
|
||||
listed = await list_mcps.on_invoke_tool(_ctx(registry), "{}")
|
||||
assert [c["id"] for c in listed["connections"]] == ["fs"]
|
||||
|
||||
await session.aclose()
|
||||
|
||||
|
||||
async def _pump_until(predicate: Callable[[], bool], *, limit: int = 100) -> None:
|
||||
"""Yield to the event loop until ``predicate`` holds, so a background
|
||||
supervising task can advance its reconnect without a real timer."""
|
||||
for _ in range(limit):
|
||||
if predicate():
|
||||
return
|
||||
await asyncio.sleep(0)
|
||||
raise AssertionError("condition not reached")
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_idle_session_death_self_heals_on_reconnect(
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
) -> None:
|
||||
# A session that dies while idle (its supervising task cancelled between calls,
|
||||
# modeling the transport scope dying with no call in flight) reconnects once on
|
||||
# its own and keeps serving, rather than staying dead until a later call would
|
||||
# have triggered a reconnect.
|
||||
first = FakeMCPServer("fs", [_mcp_tool("read_file")])
|
||||
second = FakeMCPServer("fs", [_mcp_tool("read_file")])
|
||||
built = iter([first, second])
|
||||
monkeypatch.setattr(mcp_client, "_build_server", lambda _config: next(built))
|
||||
|
||||
session = await _started_session(_secret_config("fs"))
|
||||
registry = McpRegistry()
|
||||
registry.add(name="fs", session=session, tool_count=1)
|
||||
|
||||
assert session._task is not None
|
||||
session._task.cancel() # idle transport death: no call in flight
|
||||
await _pump_until(lambda: session.server is second)
|
||||
assert session.is_dead is False
|
||||
|
||||
# The reconnected session serves calls normally.
|
||||
out = await call_mcp.on_invoke_tool(
|
||||
_ctx(registry), json.dumps({"connection": "fs", "tool": "read_file"})
|
||||
)
|
||||
assert out == {"type": "text", "text": "routed:read_file"}
|
||||
|
||||
await session.aclose()
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_flapping_idle_session_is_marked_dead_without_looping(
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
) -> None:
|
||||
# If a session reconnects after an idle death but dies again before serving any
|
||||
# call, the supervisor stops reconnecting and marks the connection dead, so a
|
||||
# server that instantly drops on connect cannot spin in a reconnect loop.
|
||||
first = FakeMCPServer("fs", [_mcp_tool("read_file")])
|
||||
second = FakeMCPServer("fs", [_mcp_tool("read_file")])
|
||||
built = iter([first, second])
|
||||
monkeypatch.setattr(mcp_client, "_build_server", lambda _config: next(built))
|
||||
|
||||
session = await _started_session(_secret_config("fs"))
|
||||
registry = McpRegistry()
|
||||
registry.add(name="fs", session=session, tool_count=1)
|
||||
|
||||
assert session._task is not None
|
||||
# First idle death heals onto the second server (only two builds ever happen).
|
||||
session._task.cancel()
|
||||
await _pump_until(lambda: session.server is second)
|
||||
assert session.is_dead is False
|
||||
|
||||
# Second idle death before any call is served: give up rather than reconnect.
|
||||
session._task.cancel()
|
||||
await _pump_until(lambda: session._task is not None and session._task.done())
|
||||
assert session.is_dead is True
|
||||
|
||||
out = await call_mcp.on_invoke_tool(
|
||||
_ctx(registry), json.dumps({"connection": "fs", "tool": "read_file"})
|
||||
)
|
||||
assert out["success"] is False
|
||||
assert "unavailable" in out["content"]
|
||||
|
||||
await session.aclose()
|
||||
|
||||
|
||||
class _HangingCallServer(FakeMCPServer):
|
||||
"""A connected server whose ``call_tool`` never returns, modeling a hung
|
||||
in-flight call so teardown can be tested for boundedness."""
|
||||
|
||||
def __init__(self, name: str, tools: list[MCPTool]) -> None:
|
||||
super().__init__(name, tools)
|
||||
self.cleaned = False
|
||||
|
||||
async def call_tool(
|
||||
self,
|
||||
tool_name: str,
|
||||
arguments: dict[str, Any] | None,
|
||||
meta: dict[str, Any] | None = None,
|
||||
) -> CallToolResult:
|
||||
await asyncio.Event().wait()
|
||||
raise AssertionError("unreachable")
|
||||
|
||||
async def cleanup(self) -> None:
|
||||
self.cleaned = True
|
||||
|
||||
|
||||
class _HangingConnectServer(FakeMCPServer):
|
||||
"""A server whose ``connect`` never finishes, so ``start`` blocks on readiness
|
||||
and can be cancelled mid-connect."""
|
||||
|
||||
def __init__(self, name: str, tools: list[MCPTool]) -> None:
|
||||
super().__init__(name, tools)
|
||||
self.cleaned = False
|
||||
|
||||
async def connect(self) -> None:
|
||||
await asyncio.Event().wait()
|
||||
|
||||
async def cleanup(self) -> None:
|
||||
self.cleaned = True
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_aclose_is_bounded_when_an_in_flight_call_hangs(
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
) -> None:
|
||||
# A hung call must not queue the shutdown sentinel behind itself forever:
|
||||
# aclose falls back to cancelling the supervising task, and cleanup still runs.
|
||||
monkeypatch.setattr(mcp_session_mod, "_SHUTDOWN_TIMEOUT", 0.2)
|
||||
server = _HangingCallServer("fs", [_mcp_tool("read_file")])
|
||||
monkeypatch.setattr(mcp_client, "_build_server", lambda _config: server)
|
||||
|
||||
session = await _started_session(_secret_config("fs"))
|
||||
call = asyncio.create_task(session.dispatch("read_file", {}, label="fs_read_file"))
|
||||
await asyncio.sleep(0.05) # let the serve loop pick up the request and hang
|
||||
|
||||
# Must return promptly rather than block on the hung call.
|
||||
await asyncio.wait_for(session.aclose(), timeout=3.0)
|
||||
assert session._task is not None and session._task.done()
|
||||
assert server.cleaned is True
|
||||
|
||||
# The abandoned caller gets a value (dead), not a hang.
|
||||
out = await asyncio.wait_for(call, timeout=3.0)
|
||||
assert isinstance(out, dict) and out["success"] is False
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_aclose_cleans_up_when_connect_is_cancelled_mid_await(
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
) -> None:
|
||||
# If the scan is cancelled while start() awaits readiness, the readiness future
|
||||
# is cancelled; aclose must not raise on it and must still cancel + clean up the
|
||||
# partially connected supervisor.
|
||||
server = _HangingConnectServer("fs", [_mcp_tool("read_file")])
|
||||
monkeypatch.setattr(mcp_client, "_build_server", lambda _config: server)
|
||||
|
||||
session = SupervisedMcpSession(_secret_config("fs"))
|
||||
start = asyncio.create_task(session.start())
|
||||
await asyncio.sleep(0.05) # let the supervisor reach the hanging connect()
|
||||
start.cancel()
|
||||
with contextlib.suppress(asyncio.CancelledError):
|
||||
await start
|
||||
|
||||
await asyncio.wait_for(session.aclose(), timeout=3.0)
|
||||
assert session._task is not None and session._task.done()
|
||||
assert server.cleaned is True
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_a_session_death_is_contained_and_other_connections_survive(
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
) -> None:
|
||||
# A background failure that surfaces as a cancellation (the transport scope
|
||||
# dying) is contained to that one session: the caller gets a value, not a
|
||||
# raised CancelledError, and a second healthy connection keeps working.
|
||||
dying = _DyingHttpServer("dying", [_mcp_tool("read_file")], death=asyncio.CancelledError())
|
||||
healthy = FakeMCPServer("healthy", [_mcp_tool("read_file")])
|
||||
dying_builds = {"n": 0}
|
||||
|
||||
def _build(config: McpConnectionConfig) -> MCPServer:
|
||||
if config.name == "healthy":
|
||||
return healthy
|
||||
# The dying connection connects once, then its rebuild raises, so it ends
|
||||
# up marked dead rather than recovering.
|
||||
dying_builds["n"] += 1
|
||||
if dying_builds["n"] == 1:
|
||||
return dying
|
||||
raise ConnectionError("cannot reconnect")
|
||||
|
||||
monkeypatch.setattr(mcp_client, "_build_server", _build)
|
||||
|
||||
dying_session = await _started_session(_secret_config("dying"))
|
||||
healthy_session = await _started_session(_secret_config("healthy"))
|
||||
registry = McpRegistry()
|
||||
registry.add(name="dying", session=dying_session, tool_count=1)
|
||||
registry.add(name="healthy", session=healthy_session, tool_count=1)
|
||||
|
||||
dead_out = await call_mcp.on_invoke_tool(
|
||||
_ctx(registry), json.dumps({"connection": "dying", "tool": "read_file"})
|
||||
)
|
||||
# Contained: a value came back rather than a CancelledError tearing down the run.
|
||||
assert isinstance(dead_out, dict)
|
||||
assert dead_out["success"] is False
|
||||
|
||||
good_out = await call_mcp.on_invoke_tool(
|
||||
_ctx(registry), json.dumps({"connection": "healthy", "tool": "read_file"})
|
||||
)
|
||||
assert good_out == {"type": "text", "text": "routed:read_file"}
|
||||
|
||||
await dying_session.aclose()
|
||||
await healthy_session.aclose()
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_reconnect_reuses_the_stored_config_and_never_logs_the_token(
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
caplog: pytest.LogCaptureFixture,
|
||||
) -> None:
|
||||
# The reconnect path rebuilds from the config held on the session, reusing the
|
||||
# same bearer token, and that token never reaches a log line, a repr, or the
|
||||
# inventory list_mcps emits.
|
||||
seen_tokens: list[str | None] = []
|
||||
|
||||
def _build(config: McpConnectionConfig) -> MCPServer:
|
||||
seen_tokens.append(config.auth.token if config.auth else None)
|
||||
if len(seen_tokens) == 1:
|
||||
return _DyingHttpServer("fs", [_mcp_tool("read_file")], death=ConnectionError("403"))
|
||||
return FakeMCPServer("fs", [_mcp_tool("read_file")])
|
||||
|
||||
monkeypatch.setattr(mcp_client, "_build_server", _build)
|
||||
|
||||
config = _secret_config("fs")
|
||||
token = config.auth.token if config.auth else ""
|
||||
session = await _started_session(config)
|
||||
registry = McpRegistry()
|
||||
entry = registry.add(name="fs", session=session, tool_count=1)
|
||||
|
||||
with caplog.at_level("DEBUG", logger="strix.tools.mcp.session"):
|
||||
out = await call_mcp.on_invoke_tool(
|
||||
_ctx(registry), json.dumps({"connection": "fs", "tool": "read_file"})
|
||||
)
|
||||
|
||||
# The retry succeeded, and both the initial connect and the reconnect used the
|
||||
# same token from the stored config (never re-fetched).
|
||||
assert out == {"type": "text", "text": "routed:read_file"}
|
||||
assert seen_tokens == [token, token]
|
||||
|
||||
# The token appears in no log line, no repr of the session or entry, and not in
|
||||
# the inventory the agent sees.
|
||||
assert token not in caplog.text
|
||||
assert token not in repr(session)
|
||||
assert token not in repr(entry)
|
||||
listed = await list_mcps.on_invoke_tool(_ctx(registry), "{}")
|
||||
assert token not in json.dumps(listed)
|
||||
# The config is still reachable in memory for the reconnect path.
|
||||
assert entry.config is config
|
||||
|
||||
await session.aclose()
|
||||
|
||||
|
||||
# --- connection status signal ------------------------------------------------
|
||||
|
||||
|
||||
class _RaisingMCPServer(FakeMCPServer):
|
||||
"""A connected server whose every call raises, so the session dies."""
|
||||
|
||||
async def call_tool(
|
||||
self,
|
||||
tool_name: str,
|
||||
arguments: dict[str, Any] | None,
|
||||
meta: dict[str, Any] | None = None,
|
||||
) -> CallToolResult:
|
||||
raise RuntimeError("connection lost")
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_session_on_dead_fires_once_on_the_death_transition() -> None:
|
||||
# An adopted session with no config cannot reconnect, so the first failed
|
||||
# call marks it dead; the on-dead callback fires exactly once, on the edge.
|
||||
server = _RaisingMCPServer("db", [_mcp_tool("read")])
|
||||
session = SupervisedMcpSession.adopt(server, name="db")
|
||||
fires: list[int] = []
|
||||
session.set_on_dead(lambda: fires.append(1))
|
||||
|
||||
out = await session.dispatch("read", {}, label="db_read")
|
||||
|
||||
assert session.is_dead is True
|
||||
assert isinstance(out, dict) and out.get("success") is False
|
||||
assert fires == [1]
|
||||
|
||||
# A later call to the already-dead session must not fire the callback again.
|
||||
await session.dispatch("read", {}, label="db_read")
|
||||
assert fires == [1]
|
||||
|
||||
|
||||
def test_registry_statuses_report_the_live_dead_flag_and_provider() -> None:
|
||||
registry = McpRegistry()
|
||||
alive = SupervisedMcpSession.adopt(FakeMCPServer("a", []), name="a")
|
||||
gone = SupervisedMcpSession.adopt(FakeMCPServer("b", []), name="b")
|
||||
registry.add(name="a", session=alive, tool_count=2, provider="supabase")
|
||||
registry.add(name="b", session=gone, tool_count=1, provider=None)
|
||||
gone._mark_dead()
|
||||
|
||||
statuses = {status.name: status for status in registry.statuses()}
|
||||
assert statuses["a"].dead is False
|
||||
assert statuses["a"].tool_count == 2
|
||||
assert statuses["a"].provider == "supabase"
|
||||
assert statuses["b"].dead is True
|
||||
assert statuses["b"].provider is None
|
||||
|
|
|
|||
|
|
@ -63,6 +63,34 @@ def report_state(tmp_path: Path, monkeypatch: pytest.MonkeyPatch) -> ReportState
|
|||
return state
|
||||
|
||||
|
||||
def test_record_mcp_connection_status_persists_and_dedupes(
|
||||
report_state: ReportState, monkeypatch: pytest.MonkeyPatch
|
||||
) -> None:
|
||||
"""The roster lands on the run record so run.json carries it for the viewer,
|
||||
and an unchanged re-write is a no-op (it does not re-save)."""
|
||||
roster = [{"name": "local_fs", "provider": None, "tool_count": 3, "dead": False}]
|
||||
report_state.record_mcp_connection_status(roster)
|
||||
assert report_state.run_record["mcp_connection_status"] == roster
|
||||
|
||||
saves = 0
|
||||
original_save = report_state.save_run_data
|
||||
|
||||
def _counting_save(*args: Any, **kwargs: Any) -> None:
|
||||
nonlocal saves
|
||||
saves += 1
|
||||
original_save(*args, **kwargs)
|
||||
|
||||
monkeypatch.setattr(report_state, "save_run_data", _counting_save)
|
||||
report_state.record_mcp_connection_status(roster)
|
||||
assert saves == 0, "an identical roster must not trigger another save"
|
||||
|
||||
report_state.record_mcp_connection_status(
|
||||
[{"name": "local_fs", "provider": None, "tool_count": 3, "dead": True}]
|
||||
)
|
||||
assert saves == 1
|
||||
assert report_state.run_record["mcp_connection_status"][0]["dead"] is True
|
||||
|
||||
|
||||
async def test_create_report_persists_new_fields(report_state: ReportState) -> None:
|
||||
result = await _do_create(
|
||||
title="Reflected XSS in search",
|
||||
|
|
|
|||
|
|
@ -140,3 +140,51 @@ async def test_supplied_requests_are_attached_and_the_user_file_is_not_read(
|
|||
)
|
||||
|
||||
assert captured == [supplied]
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_roster_is_persisted_even_without_a_status_sink(
|
||||
monkeypatch: pytest.MonkeyPatch, tmp_path: Any
|
||||
) -> None:
|
||||
"""The viewer reads the roster off disk, so persistence must not depend on the
|
||||
interface status sink: with ``mcp_status_sink=None`` the connect-time roster is
|
||||
still written, carrying only the non-secret name/provider/tool_count/dead."""
|
||||
_wire_runner(monkeypatch, tmp_path)
|
||||
monkeypatch.setattr(
|
||||
mcp_pkg,
|
||||
"load_user_mcp_configs",
|
||||
lambda: [McpConnectionConfig(name="local_fs", transport="stdio", command="npx")],
|
||||
)
|
||||
|
||||
class _FakeSession:
|
||||
is_dead = False
|
||||
|
||||
def set_on_dead(self, _callback: Any) -> None:
|
||||
return None
|
||||
|
||||
async def _attach(_requests: list[McpConnectionRequest], registry: Any) -> list[Any]:
|
||||
registry.add(name="local_fs", session=_FakeSession(), tool_count=3, provider=None)
|
||||
entry = registry.get("local_fs")
|
||||
return [types.SimpleNamespace(name="local_fs", tool_count=3, session=entry.session)]
|
||||
|
||||
monkeypatch.setattr(mcp_pkg, "attach_mcp_requests", _attach)
|
||||
|
||||
persisted: list[list[dict[str, Any]]] = []
|
||||
|
||||
def _capture_persist(roster: list[dict[str, Any]]) -> None:
|
||||
persisted.append(roster)
|
||||
|
||||
monkeypatch.setattr(runner, "_persist_mcp_status", _capture_persist)
|
||||
|
||||
await runner.run_strix_scan(
|
||||
scan_config={"targets": [], "scan_mode": "deep"},
|
||||
scan_id="scan-persist",
|
||||
image="img",
|
||||
coordinator=AgentCoordinator(),
|
||||
mcp_status_sink=None,
|
||||
)
|
||||
|
||||
assert persisted, "roster must persist even when no status sink is attached"
|
||||
assert persisted[-1] == [
|
||||
{"name": "local_fs", "provider": None, "tool_count": 3, "dead": False}
|
||||
]
|
||||
|
|
|
|||
|
|
@ -194,9 +194,13 @@ async def test_mcp_available_flag_set_when_a_connection_attaches(
|
|||
scope_context: dict[str, Any] = {"scope": "built-in"}
|
||||
captured = _patch_engine_scaffold(monkeypatch, tmp_path, scope_context)
|
||||
|
||||
async def _aclose() -> None:
|
||||
return None
|
||||
|
||||
async def _attach(_requests: Any, registry: Any) -> list[Any]:
|
||||
registry.add(name="fs", server=object(), purpose="local files", tool_count=2)
|
||||
return [types.SimpleNamespace(name="fs", tool_count=2, server=object())]
|
||||
session = types.SimpleNamespace(aclose=_aclose)
|
||||
return [types.SimpleNamespace(name="fs", tool_count=2, session=session)]
|
||||
|
||||
monkeypatch.setattr(mcp_pkg, "attach_mcp_requests", _attach)
|
||||
|
||||
|
|
|
|||
|
|
@ -61,6 +61,25 @@ async def test_setup_state_is_serializable() -> None:
|
|||
assert snapshot["diff_base"] is None
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_connections_snapshot_reflects_the_pushed_mcp_roster() -> None:
|
||||
controller = TuiController(args())
|
||||
# A run with no MCP connections carries an empty roster, so the sidebar
|
||||
# omits the panel entirely.
|
||||
assert controller.snapshot()["connections"] == []
|
||||
|
||||
controller.set_mcp_connections(
|
||||
[
|
||||
{"name": "supabase", "tool_count": 3, "dead": False},
|
||||
{"name": "vercel", "tool_count": 1, "dead": True},
|
||||
]
|
||||
)
|
||||
assert controller.snapshot()["connections"] == [
|
||||
{"name": "supabase", "tool_count": 3, "dead": False},
|
||||
{"name": "vercel", "tool_count": 1, "dead": True},
|
||||
]
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_setup_instruction_starts_from_cli_and_can_be_cleared() -> None:
|
||||
setup_args = args()
|
||||
|
|
|
|||
|
|
@ -96,6 +96,25 @@ def test_read_run_summary_finished_flag(tmp_path: Path) -> None:
|
|||
assert read_run_summary(partial)["finished"] is False
|
||||
|
||||
|
||||
def test_read_run_summary_surfaces_mcp_connection_status(tmp_path: Path) -> None:
|
||||
"""The engine persists the non-secret MCP roster under mcp_connection_status;
|
||||
read_run_summary spreads the whole record, so /api/run carries it to the
|
||||
viewer verbatim."""
|
||||
run_dir = _make_run(tmp_path, "mcp", status="running", end_time=None)
|
||||
roster = [
|
||||
{"name": "local_fs", "provider": None, "tool_count": 3, "dead": False},
|
||||
{"name": "db", "provider": "supabase", "tool_count": 7, "dead": True},
|
||||
]
|
||||
record = {
|
||||
"run_name": "mcp",
|
||||
"status": "running",
|
||||
"end_time": None,
|
||||
"mcp_connection_status": roster,
|
||||
}
|
||||
(run_dir / "run.json").write_text(json.dumps(record), encoding="utf-8")
|
||||
assert read_run_summary(run_dir)["mcp_connection_status"] == roster
|
||||
|
||||
|
||||
def test_read_missing_artifacts_return_defaults(tmp_path: Path) -> None:
|
||||
run_dir = _make_run(tmp_path, "empty", status="running", end_time=None)
|
||||
assert read_vulnerabilities(run_dir) == []
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue