diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml index 0027022..c9ba380 100644 --- a/.github/workflows/ci.yml +++ b/.github/workflows/ci.yml @@ -99,4 +99,6 @@ jobs: run: make build-all - name: Verify binary - run: ./bin/rootly version + run: | + ./bin/rootly version + ./bin/rootly --version diff --git a/go.mod b/go.mod index e6858a9..82c9975 100644 --- a/go.mod +++ b/go.mod @@ -3,39 +3,74 @@ module github.com/rootlyhq/rootly-cli go 1.26.10 require ( + charm.land/bubbles/v2 v2.2.1 + charm.land/bubbletea/v2 v2.1.0 + charm.land/glamour/v2 v2.0.1 + charm.land/lipgloss/v2 v2.0.6 + github.com/charmbracelet/colorprofile v0.4.3 + github.com/charmbracelet/x/ansi v0.11.9 + github.com/google/uuid v1.6.0 + github.com/guptarohit/asciigraph v0.10.0 github.com/jedib0t/go-pretty/v6 v6.8.3 github.com/mattn/go-isatty v0.0.24 - github.com/rootlyhq/rootly-go v0.14.0 + github.com/modelcontextprotocol/go-sdk v1.8.0 + github.com/oapi-codegen/nullable v1.2.0 + github.com/oapi-codegen/runtime v1.8.0 + github.com/pelletier/go-toml/v2 v2.2.4 + github.com/rootlyhq/rootly-go v0.16.0 github.com/spf13/cobra v1.10.2 + github.com/spf13/pflag v1.0.10 github.com/spf13/viper v1.21.0 + github.com/yuin/goldmark v1.7.8 + go.starlark.net v0.0.0-20261021141710-4f4267fb140d golang.org/x/oauth2 v0.37.0 gopkg.in/yaml.v3 v3.0.1 + mvdan.cc/sh/v3 v3.14.1 ) require ( + github.com/alecthomas/chroma/v2 v2.14.0 // indirect github.com/apapsch/go-jsonmerge/v2 v2.0.0 // indirect + github.com/atotto/clipboard v0.1.4 // indirect + github.com/aymerick/douceur v0.2.0 // indirect + github.com/charmbracelet/ultraviolet v0.0.0-20261008173134-6b8d4baf91b4 // indirect + github.com/charmbracelet/x/exp/slice v0.0.0-20250327172914-2fdc97757edf // indirect + github.com/charmbracelet/x/term v0.2.2 // indirect + github.com/charmbracelet/x/termios v0.1.1 // indirect + github.com/charmbracelet/x/windows v0.2.2 // indirect + github.com/clipperhouse/displaywidth v0.11.0 // indirect github.com/clipperhouse/uax29/v2 v2.7.0 // indirect + github.com/dlclark/regexp2 v1.11.0 // indirect github.com/fsnotify/fsnotify v1.9.0 // indirect - github.com/getkin/kin-openapi v0.149.0 // indirect + github.com/getkin/kin-openapi v0.152.0 // indirect github.com/go-openapi/jsonpointer v0.22.5 // indirect github.com/go-openapi/swag/jsonname v0.25.5 // indirect github.com/go-viper/mapstructure/v2 v2.5.0 // indirect - github.com/google/go-cmp v0.7.0 // indirect - github.com/google/uuid v1.6.0 // indirect + github.com/google/jsonschema-go v0.4.3 // indirect + github.com/gorilla/css v1.0.1 // indirect github.com/inconshreveable/mousetrap v1.1.0 // indirect - github.com/mattn/go-runewidth v0.0.20 // indirect - github.com/oapi-codegen/nullable v1.2.0 // indirect - github.com/oapi-codegen/runtime v1.7.0 // indirect + github.com/lucasb-eyer/go-colorful v1.4.1 // indirect + github.com/mattn/go-runewidth v0.0.30 // indirect + github.com/microcosm-cc/bluemonday v1.0.27 // indirect + github.com/muesli/cancelreader v0.2.2 // indirect github.com/oasdiff/yaml v0.1.1 // indirect github.com/oasdiff/yaml3 v0.0.14 // indirect - github.com/pelletier/go-toml/v2 v2.2.4 // indirect + github.com/rivo/uniseg v0.4.7 // indirect github.com/sagikazarmark/locafero v0.12.0 // indirect + github.com/sahilm/fuzzy v0.1.3 // indirect github.com/santhosh-tekuri/jsonschema/v6 v6.0.3 // indirect + github.com/segmentio/asm v1.1.3 // indirect + github.com/segmentio/encoding v0.5.4 // indirect github.com/spf13/afero v1.15.0 // indirect github.com/spf13/cast v1.10.0 // indirect - github.com/spf13/pflag v1.0.10 // indirect github.com/subosito/gotenv v1.6.0 // indirect + github.com/xo/terminfo v1.0.0 // indirect + github.com/yosida95/uritemplate/v3 v3.0.2 // indirect + github.com/yuin/goldmark-emoji v1.0.5 // indirect go.yaml.in/yaml/v3 v3.0.4 // indirect - golang.org/x/sys v0.41.0 // indirect + golang.org/x/net v0.48.0 // indirect + golang.org/x/sync v0.22.0 // indirect + golang.org/x/sys v0.47.0 // indirect golang.org/x/text v0.34.0 // indirect + golang.org/x/time v0.15.0 // indirect ) diff --git a/go.sum b/go.sum index 9011d79..dc66f9c 100644 --- a/go.sum +++ b/go.sum @@ -1,7 +1,47 @@ +charm.land/bubbles/v2 v2.2.1 h1:Fq1+qm5hV6GkvzLQDhCBpXXE5tLgvh1PRriCLwSvIQU= +charm.land/bubbles/v2 v2.2.1/go.mod h1:wdMgn+sje1KNXdwFizIWjbf328fIUBxqEmJ/vYPo8yc= +charm.land/bubbletea/v2 v2.1.0 h1:A9fj9h5LFrvQhQHFNggp8K8rvY6CUfgn9dMbb6p2HIw= +charm.land/bubbletea/v2 v2.1.0/go.mod h1:xWHeyxIX9/wiFU2kgPW04ojOZyTzrWf+8lGjUCaftiU= +charm.land/glamour/v2 v2.0.1 h1:xl+r00A4aJWU0z8fgwKd9fQQ4rsphqGUzuEiXZP5n+c= +charm.land/glamour/v2 v2.0.1/go.mod h1:jo9z8XqVKPeEFMVdvCRLGk++RyJ3CdUwgNr7EvXLw3k= +charm.land/lipgloss/v2 v2.0.6 h1:EaGKeuA8FvF+v2BT5VmZd2LoYLaMZJXA5n34th8nCIQ= +charm.land/lipgloss/v2 v2.0.6/go.mod h1:ipDDJNSGa1hlwDtSfW1s2/xR8Vdhbut4PXh2zEKZd0Q= +github.com/MakeNowJust/heredoc v1.0.0 h1:cXCdzVdstXyiTqTvfqk9SDHpKNjxuom+DOlyEeQ4pzQ= +github.com/MakeNowJust/heredoc v1.0.0/go.mod h1:mG5amYoWBHf8vpLOuehzbGGw0EHxpZZ6lCpQ4fNJ8LE= github.com/RaveNoX/go-jsoncommentstrip v1.0.0/go.mod h1:78ihd09MekBnJnxpICcwzCMzGrKSKYe4AqU6PDYYpjk= +github.com/alecthomas/assert/v2 v2.7.0 h1:QtqSACNS3tF7oasA8CU6A6sXZSBDqnm7RfpLl9bZqbE= +github.com/alecthomas/assert/v2 v2.7.0/go.mod h1:Bze95FyfUr7x34QZrjL+XP+0qgp/zg8yS+TtBj1WA3k= +github.com/alecthomas/chroma/v2 v2.14.0 h1:R3+wzpnUArGcQz7fCETQBzO5n9IMNi13iIs46aU4V9E= +github.com/alecthomas/chroma/v2 v2.14.0/go.mod h1:QolEbTfmUHIMVpBqxeDnNBj2uoeI4EbYP4i6n68SG4I= +github.com/alecthomas/repr v0.4.0 h1:GhI2A8MACjfegCPVq9f1FLvIBS+DrQ2KQBFZP1iFzXc= +github.com/alecthomas/repr v0.4.0/go.mod h1:Fr0507jx4eOXV7AlPV6AVZLYrLIuIeSOWtW57eE/O/4= github.com/apapsch/go-jsonmerge/v2 v2.0.0 h1:axGnT1gRIfimI7gJifB699GoE/oq+F2MU7Dml6nw9rQ= github.com/apapsch/go-jsonmerge/v2 v2.0.0/go.mod h1:lvDnEdqiQrp0O42VQGgmlKpxL1AP2+08jFMw88y4klk= +github.com/atotto/clipboard v0.1.4 h1:EH0zSVneZPSuFR11BlR9YppQTVDbh5+16AmcJi4g1z4= +github.com/atotto/clipboard v0.1.4/go.mod h1:ZY9tmq7sm5xIbd9bOK4onWV4S6X0u6GY7Vn0Yu86PYI= +github.com/aymanbagabas/go-udiff v0.4.1 h1:OEIrQ8maEeDBXQDoGCbbTTXYJMYRCRO1fnodZ12Gv5o= +github.com/aymanbagabas/go-udiff v0.4.1/go.mod h1:0L9PGwj20lrtmEMeyw4WKJ/TMyDtvAoK9bf2u/mNo3w= +github.com/aymerick/douceur v0.2.0 h1:Mv+mAeH1Q+n9Fr+oyamOlAkUNPWPlA8PPGR0QAaYuPk= +github.com/aymerick/douceur v0.2.0/go.mod h1:wlT5vV2O3h55X9m7iVYN0TBM0NH/MmbLnd30/FjWUq4= github.com/bmatcuk/doublestar v1.1.1/go.mod h1:UD6OnuiIn0yFxxA2le/rnRU1G4RaI4UvFv1sNto9p6w= +github.com/charmbracelet/colorprofile v0.4.3 h1:QPa1IWkYI+AOB+fE+mg/5/4HRMZcaXex9t5KX76i20Q= +github.com/charmbracelet/colorprofile v0.4.3/go.mod h1:/zT4BhpD5aGFpqQQqw7a+VtHCzu+zrQtt1zhMt9mR4Q= +github.com/charmbracelet/ultraviolet v0.0.0-20261008173134-6b8d4baf91b4 h1:hISLzUbwcvJbGoNR6A3L/TpguRRnACBmZOee8LcHLJw= +github.com/charmbracelet/ultraviolet v0.0.0-20261008173134-6b8d4baf91b4/go.mod h1:aJDVk+WOVM1K04rdc/71PMPoUGfjxjhvXfP3Pjz1Zio= +github.com/charmbracelet/x/ansi v0.11.9 h1:kIrqLBxtTVbLyyHO2Evoi93EjQo1mylWmpNEsfPxXAM= +github.com/charmbracelet/x/ansi v0.11.9/go.mod h1:DG/hpOnaI2QJ/KMMqLL2h0KVhluuJLazf8973P+Vy5w= +github.com/charmbracelet/x/exp/golden v0.0.0-20250806222409-83e3a29d542f h1:pk6gmGpCE7F3FcjaOEKYriCvpmIN4+6OS/RD0vm4uIA= +github.com/charmbracelet/x/exp/golden v0.0.0-20250806222409-83e3a29d542f/go.mod h1:IfZAMTHB6XkZSeXUqriemErjAWCCzT0LwjKFYCZyw0I= +github.com/charmbracelet/x/exp/slice v0.0.0-20250327172914-2fdc97757edf h1:rLG0Yb6MQSDKdB52aGX55JT1oi0P0Kuaj7wi1bLUpnI= +github.com/charmbracelet/x/exp/slice v0.0.0-20250327172914-2fdc97757edf/go.mod h1:B3UgsnsBZS/eX42BlaNiJkD1pPOUa+oF1IYC6Yd2CEU= +github.com/charmbracelet/x/term v0.2.2 h1:xVRT/S2ZcKdhhOuSP4t5cLi5o+JxklsoEObBSgfgZRk= +github.com/charmbracelet/x/term v0.2.2/go.mod h1:kF8CY5RddLWrsgVwpw4kAa6TESp6EB5y3uxGLeCqzAI= +github.com/charmbracelet/x/termios v0.1.1 h1:o3Q2bT8eqzGnGPOYheoYS8eEleT5ZVNYNy8JawjaNZY= +github.com/charmbracelet/x/termios v0.1.1/go.mod h1:rB7fnv1TgOPOyyKRJ9o+AsTU/vK5WHJ2ivHeut/Pcwo= +github.com/charmbracelet/x/windows v0.2.2 h1:IofanmuvaxnKHuV04sC0eBy/smG6kIKrWG2/jYn2GuM= +github.com/charmbracelet/x/windows v0.2.2/go.mod h1:/8XtdKZzedat74NQFn0NGlGL4soHB0YQZrETF96h75k= +github.com/clipperhouse/displaywidth v0.11.0 h1:lBc6kY44VFw+TDx4I8opi/EtL9m20WSEFgwIwO+UVM8= +github.com/clipperhouse/displaywidth v0.11.0/go.mod h1:bkrFNkf81G8HyVqmKGxsPufD3JhNl3dSqnGhOoSD/o0= github.com/clipperhouse/uax29/v2 v2.7.0 h1:+gs4oBZ2gPfVrKPthwbMzWZDaAFPGYK72F0NJv2v7Vk= github.com/clipperhouse/uax29/v2 v2.7.0/go.mod h1:EFJ2TJMRUaplDxHKj1qAEhCtQPW2tJSwu5BF98AuoVM= github.com/cpuguy83/go-md2man/v2 v2.0.6/go.mod h1:oOW0eioCTA6cOiMLiUPZOpcVxMig6NIQQ7OS05n1F4g= @@ -14,20 +54,32 @@ github.com/frankban/quicktest v1.14.6 h1:7Xjx+VpznH+oBnejlPUj8oUpdxnVs4f8XU8WnHk github.com/frankban/quicktest v1.14.6/go.mod h1:4ptaffx2x8+WTWXmUCuVU6aPUX1/Mz7zb5vbUoiM6w0= github.com/fsnotify/fsnotify v1.9.0 h1:2Ml+OJNzbYCTzsxtv8vKSFD9PbJjmhYF14k/jKC7S9k= github.com/fsnotify/fsnotify v1.9.0/go.mod h1:8jBTzvmWwFyi3Pb8djgCCO5IBqzKJ/Jwo8TRcHyHii0= -github.com/getkin/kin-openapi v0.149.0 h1:ZbhmVJ4yq5RZDUsyP8lcBcGMsjsaTqXEFt6isdtMDfA= -github.com/getkin/kin-openapi v0.149.0/go.mod h1:1+BHDzstro+P5CKtPy1X4PfofnFgmRe6uvMy9+r9fKY= +github.com/getkin/kin-openapi v0.152.0 h1:XRVo77lMRyOGMHf1paMu9lxX3wzuQDZeprssZJnz/x8= +github.com/getkin/kin-openapi v0.152.0/go.mod h1:1+BHDzstro+P5CKtPy1X4PfofnFgmRe6uvMy9+r9fKY= github.com/go-openapi/jsonpointer v0.22.5 h1:8on/0Yp4uTb9f4XvTrM2+1CPrV05QPZXu+rvu2o9jcA= github.com/go-openapi/jsonpointer v0.22.5/go.mod h1:gyUR3sCvGSWchA2sUBJGluYMbe1zazrYWIkWPjjMUY0= github.com/go-openapi/swag/jsonname v0.25.5 h1:8p150i44rv/Drip4vWI3kGi9+4W9TdI3US3uUYSFhSo= github.com/go-openapi/swag/jsonname v0.25.5/go.mod h1:jNqqikyiAK56uS7n8sLkdaNY/uq6+D2m2LANat09pKU= github.com/go-openapi/testify/v2 v2.4.0 h1:8nsPrHVCWkQ4p8h1EsRVymA2XABB4OT40gcvAu+voFM= github.com/go-openapi/testify/v2 v2.4.0/go.mod h1:HCPmvFFnheKK2BuwSA0TbbdxJ3I16pjwMkYkP4Ywn54= +github.com/go-quicktest/qt v1.102.0 h1:HSQxCeh5YZH3EL3W39ixjtyaEhcWSXQHtHnMBzSs474= +github.com/go-quicktest/qt v1.102.0/go.mod h1:p4lGIVX+8Wa6ZPNDvqcxq36XpUDLh42FLetFU7odllI= github.com/go-viper/mapstructure/v2 v2.5.0 h1:vM5IJoUAy3d7zRSVtIwQgBj7BiWtMPfmPEgAXnvj1Ro= github.com/go-viper/mapstructure/v2 v2.5.0/go.mod h1:oJDH3BJKyqBA2TXFhDsKDGDTlndYOZ6rGS0BRZIxGhM= +github.com/golang-jwt/jwt/v5 v5.3.1 h1:kYf81DTWFe7t+1VvL7eS+jKFVWaUnK9cB1qbwn63YCY= +github.com/golang-jwt/jwt/v5 v5.3.1/go.mod h1:fxCRLWMO43lRc8nhHWY6LGqRcf+1gQWArsqaEUEa5bE= github.com/google/go-cmp v0.7.0 h1:wk8382ETsv4JYUZwIsn6YpYiWiBsYLSJiTsyBybVuN8= github.com/google/go-cmp v0.7.0/go.mod h1:pXiqmnSA92OHEEa9HXL2W4E7lf9JzCmGVUdgjX3N/iU= +github.com/google/jsonschema-go v0.4.3 h1:/DBOLZTfDow7pe2GmaJNhltueGTtDKICi8V8p+DQPd0= +github.com/google/jsonschema-go v0.4.3/go.mod h1:r5quNTdLOYEz95Ru18zA0ydNbBuYoo9tgaYcxEYhJVE= github.com/google/uuid v1.6.0 h1:NIvaJDMOsjHA8n1jAhLSgzrAzy1Hgr+hNrb57e+94F0= github.com/google/uuid v1.6.0/go.mod h1:TIyPZe4MgqvfeYDBFedMoGGpEw/LqOeaOT+nhxU+yHo= +github.com/gorilla/css v1.0.1 h1:ntNaBIghp6JmvWnxbZKANoLyuXTPZ4cAMlo6RyhlbO8= +github.com/gorilla/css v1.0.1/go.mod h1:BvnYkspnSzMmwRK+b8/xgNPLiIuNZr6vbZBTPQ2A3b0= +github.com/guptarohit/asciigraph v0.10.0 h1:LmbFXSHZOhaQxjJYexdRk7TzoC5sJ7vDTEjP1YUbKgY= +github.com/guptarohit/asciigraph v0.10.0/go.mod h1:dYl5wwK4gNsnFf9Zp+l06rFiDZ5YtXM6x7SRWZ3KGag= +github.com/hexops/gotextdiff v1.0.3 h1:gitA9+qJrrTCsiCl7+kh75nPqQt1cx4ZkudSTLoUqJM= +github.com/hexops/gotextdiff v1.0.3/go.mod h1:pSWU5MAI3yDq+fZBTazCSJysOMbxWL1BSow5/V2vxeg= github.com/inconshreveable/mousetrap v1.1.0 h1:wN+x4NVGpMsO7ErUn/mUI3vEoE6Jt13X2s0bqwp9tc8= github.com/inconshreveable/mousetrap v1.1.0/go.mod h1:vpF70FUmC8bwa3OWnCshd2FqLfsEA9PFc4w1p2J65bw= github.com/jedib0t/go-pretty/v6 v6.8.3 h1:yVSk5aemoYHCvcrtqyXklwqcgHQIQzmy/oUzFlmffSQ= @@ -37,14 +89,24 @@ github.com/kr/pretty v0.3.1 h1:flRD4NNwYAUpkphVc1HcthR4KEIFJ65n8Mw5qdRn3LE= github.com/kr/pretty v0.3.1/go.mod h1:hoEshYVHaxMs3cyo3Yncou5ZscifuDolrwPKZanG3xk= github.com/kr/text v0.2.0 h1:5Nx0Ya0ZqY2ygV366QzturHI13Jq95ApcVaJBhpS+AY= github.com/kr/text v0.2.0/go.mod h1:eLer722TekiGuMkidMxC/pM04lWEeraHUUmBw8l2grE= +github.com/kylelemons/godebug v1.1.0 h1:RPNrshWIDI6G2gRW9EHilWtl7Z6Sb1BR0xunSBf0SNc= +github.com/kylelemons/godebug v1.1.0/go.mod h1:9/0rRGxNHcop5bhtWyNeEfOS8JIWk580+fNqagV/RAw= +github.com/lucasb-eyer/go-colorful v1.4.1 h1:1EO+WB73+EH8EVbzlrG3KLAfEypQWVHIBqlTf+2hNss= +github.com/lucasb-eyer/go-colorful v1.4.1/go.mod h1:R4dSotOR9KMtayYi1e77YzuveK+i7ruzyGqttikkLy0= github.com/mattn/go-isatty v0.0.24 h1:tGZZoVgT/KiqK1c8ocVLeDS8BSWMRd47J3Lbz7vsReI= github.com/mattn/go-isatty v0.0.24/go.mod h1:nMCL3Zebbrt45jsMDgnfIwz6ydEQApk5oEI3HqDio6A= -github.com/mattn/go-runewidth v0.0.20 h1:WcT52H91ZUAwy8+HUkdM3THM6gXqXuLJi9O3rjcQQaQ= -github.com/mattn/go-runewidth v0.0.20/go.mod h1:XBkDxAl56ILZc9knddidhrOlY5R/pDhgLpndooCuJAs= +github.com/mattn/go-runewidth v0.0.30 h1:+KUuiDA4fF0R1p5FeueHefjDm+GIM+kWfFnDjybOPgk= +github.com/mattn/go-runewidth v0.0.30/go.mod h1:3qAiGCV4Koz/yuveO58qUefmUTRm8r0IGEXZ9jeHp/8= +github.com/microcosm-cc/bluemonday v1.0.27 h1:MpEUotklkwCSLeH+Qdx1VJgNqLlpY2KXwXFM08ygZfk= +github.com/microcosm-cc/bluemonday v1.0.27/go.mod h1:jFi9vgW+H7c3V0lb6nR74Ib/DIB5OBs92Dimizgw2cA= +github.com/modelcontextprotocol/go-sdk v1.8.0 h1:KIvahhYqwtbeniWVPs3TcXEA7b8jEtwfBpOTAI+Urx4= +github.com/modelcontextprotocol/go-sdk v1.8.0/go.mod h1:dL7u98E/zjJTGzEq+j30jQ8K2k1mb6LeAH4inEcSGts= +github.com/muesli/cancelreader v0.2.2 h1:3I4Kt4BQjOR54NavqnDogx/MIoWBFa0StPA8ELUXHmA= +github.com/muesli/cancelreader v0.2.2/go.mod h1:3XuTXfFS2VjM+HTLZY9Ak0l6eUKfijIfMUZ4EgX0QYo= github.com/oapi-codegen/nullable v1.2.0 h1:VflFkDW980KhBPiFF7nWSyjg+r4Obqj8lXipV0UkP5w= github.com/oapi-codegen/nullable v1.2.0/go.mod h1:KUZ3vUzkmEKY90ksAmit2+5juDIhIZhfDl+0PwOQlFY= -github.com/oapi-codegen/runtime v1.7.0 h1:t7358VYPvNbWJ9gdAkIK/smVeHpBf6yp8VTsaZsb/7k= -github.com/oapi-codegen/runtime v1.7.0/go.mod h1:GwV7hC2hviaMzj+ITfHVRESK5J2W/GefVwIND/bMGvU= +github.com/oapi-codegen/runtime v1.8.0 h1:RJzEgQSJQADEPaG1cYuSCKYai/cCVRE2xnkrqhJs/3c= +github.com/oapi-codegen/runtime v1.8.0/go.mod h1:GwV7hC2hviaMzj+ITfHVRESK5J2W/GefVwIND/bMGvU= github.com/oasdiff/yaml v0.1.1 h1:6nHx+pn9gBRM6YpBlFZFQGCCd1nuvqOBtTD3KKTgGxY= github.com/oasdiff/yaml v0.1.1/go.mod h1:EYJNoyktvWMJ0Hmhx+6qTaqMOsalUaRGT8Sj1hNcegU= github.com/oasdiff/yaml3 v0.0.14 h1:aLJee3hxBK2H5wdXd9iPcIXb93Nty1Ge0pT171eHtkw= @@ -53,15 +115,23 @@ github.com/pelletier/go-toml/v2 v2.2.4 h1:mye9XuhQ6gvn5h28+VilKrrPoQVanw5PMw/TB0 github.com/pelletier/go-toml/v2 v2.2.4/go.mod h1:2gIqNv+qfxSVS7cM2xJQKtLSTLUE9V8t9Stt+h56mCY= github.com/pmezard/go-difflib v1.0.0 h1:4DBwDE0NGyQoBHbLQYPwSUPoCMWR5BEzIk/f1lZbAQM= github.com/pmezard/go-difflib v1.0.0/go.mod h1:iKH77koFhYxTK1pcRnkKkqfTogsbg7gZNVY4sRDYZ/4= -github.com/rogpeppe/go-internal v1.9.0 h1:73kH8U+JUqXU8lRuOHeVHaa/SZPifC7BkcraZVejAe8= -github.com/rogpeppe/go-internal v1.9.0/go.mod h1:WtVeX8xhTBvf0smdhujwtBcq4Qrzq/fJaraNFVN+nFs= -github.com/rootlyhq/rootly-go v0.14.0 h1:7LPXTkErGmZJpnb+eSL3nF5VKv8NiZ8CNhN2OC9ekwI= -github.com/rootlyhq/rootly-go v0.14.0/go.mod h1:kVBjbv7kXeDII6N8JX0MOdDjDMF6xJVLMcPzkrxjQ20= +github.com/rivo/uniseg v0.4.7 h1:WUdvkW8uEhrYfLC4ZzdpI2ztxP1I582+49Oc5Mq64VQ= +github.com/rivo/uniseg v0.4.7/go.mod h1:FN3SvrM+Zdj16jyLfmOkMNblXMcoc8DfTHruCPUcx88= +github.com/rogpeppe/go-internal v1.15.0 h1:D0RCU5rMAp+SpgkiNdrjfJ+LX4J1M32V2NeCY7EJ6hc= +github.com/rogpeppe/go-internal v1.15.0/go.mod h1:DrUVZyrJU+txYW5/1kwtXQSMFio52ZOxX7yM1VHvnxs= +github.com/rootlyhq/rootly-go v0.16.0 h1:TJ7MkmrU0KexU5OFp3DZFWfj3y8nLOlem+nvBbNdHS0= +github.com/rootlyhq/rootly-go v0.16.0/go.mod h1:/eWs3ZcY3uJLaHeTmQ9ILegtQeywXpez49RbbdPiim0= github.com/russross/blackfriday/v2 v2.1.0/go.mod h1:+Rmxgy9KzJVeS9/2gXHxylqXiyQDYRxCVz55jmeOWTM= github.com/sagikazarmark/locafero v0.12.0 h1:/NQhBAkUb4+fH1jivKHWusDYFjMOOKU88eegjfxfHb4= github.com/sagikazarmark/locafero v0.12.0/go.mod h1:sZh36u/YSZ918v0Io+U9ogLYQJ9tLLBmM4eneO6WwsI= +github.com/sahilm/fuzzy v0.1.3 h1:juByESSS32nVD81vr6tHmKmA/8zde7gE+x5CLxrzXPU= +github.com/sahilm/fuzzy v0.1.3/go.mod h1:au6//VbVSqu6DFrkL2CfjlJ5iURpNCPeE+1GwY3XsT8= github.com/santhosh-tekuri/jsonschema/v6 v6.0.3 h1:1EYB5IzjZawrrnELUi78f9fPu57HuXjmddZPjrls/28= github.com/santhosh-tekuri/jsonschema/v6 v6.0.3/go.mod h1:JXeL+ps8p7/KNMjDQk3TCwPpBy0wYklyWTfbkIzdIFU= +github.com/segmentio/asm v1.1.3 h1:WM03sfUOENvvKexOLp+pCqgb/WDjsi7EK8gIsICtzhc= +github.com/segmentio/asm v1.1.3/go.mod h1:Ld3L4ZXGNcSLRg4JBsZ3//1+f/TjYl0Mzen/DQy1EJg= +github.com/segmentio/encoding v0.5.4 h1:OW1VRern8Nw6ITAtwSZ7Idrl3MXCFwXHPgqESYfvNt0= +github.com/segmentio/encoding v0.5.4/go.mod h1:HS1ZKa3kSN32ZHVZ7ZLPLXWvOVIiZtyJnO1gPH1sKt0= github.com/spf13/afero v1.15.0 h1:b/YBCLWAJdFWJTN9cLhiXXcD7mzKn9Dm86dNnfyQw1I= github.com/spf13/afero v1.15.0/go.mod h1:NC2ByUVxtQs4b3sIUphxK0NioZnmxgyCrfzeuq8lxMg= github.com/spf13/cast v1.10.0 h1:h2x0u2shc1QuLHfxi+cTJvs30+ZAHOGRic8uyGTDWxY= @@ -80,16 +150,41 @@ github.com/stretchr/testify v1.11.1 h1:7s2iGBzp5EwR7/aIZr8ao5+dra3wiQyKjjFuvgVKu github.com/stretchr/testify v1.11.1/go.mod h1:wZwfW3scLgRK+23gO65QZefKpKQRnfz6sD981Nm4B6U= github.com/subosito/gotenv v1.6.0 h1:9NlTDc1FTs4qu0DDq7AEtTPNw6SVm7uBMsUCUjABIf8= github.com/subosito/gotenv v1.6.0/go.mod h1:Dk4QP5c2W3ibzajGcXpNraDfq2IrhjMIvMSWPKKo0FU= +github.com/xo/terminfo v1.0.0 h1:2ZpYzqWzyyytjk3TP6aJVDhkMAkc99/1xKQdA3TDTBY= +github.com/xo/terminfo v1.0.0/go.mod h1:RbqR21r5mrJuqunuUZ/Dhy/avygyECGrLceyNeo4LiM= +github.com/yosida95/uritemplate/v3 v3.0.2 h1:Ed3Oyj9yrmi9087+NczuL5BwkIc4wvTb5zIM+UJPGz4= +github.com/yosida95/uritemplate/v3 v3.0.2/go.mod h1:ILOh0sOhIJR3+L/8afwt/kE++YT040gmv5BQTMR2HP4= +github.com/yuin/goldmark v1.7.1/go.mod h1:uzxRWxtg69N339t3louHJ7+O03ezfj6PlliRlaOzY1E= +github.com/yuin/goldmark v1.7.8 h1:iERMLn0/QJeHFhxSt3p6PeN9mGnvIKSpG9YYorDMnic= +github.com/yuin/goldmark v1.7.8/go.mod h1:uzxRWxtg69N339t3louHJ7+O03ezfj6PlliRlaOzY1E= +github.com/yuin/goldmark-emoji v1.0.5 h1:EMVWyCGPlXJfUXBXpuMu+ii3TIaxbVBnEX9uaDC4cIk= +github.com/yuin/goldmark-emoji v1.0.5/go.mod h1:tTkZEbwu5wkPmgTcitqddVxY9osFZiavD+r4AzQrh1U= +go.starlark.net v0.0.0-20261021141710-4f4267fb140d h1:wmJ4xTYwnBup3nwA+3OtmpaCwd90YXVyyCHZQoiD0+0= +go.starlark.net v0.0.0-20261021141710-4f4267fb140d/go.mod h1:Iue6g6iirlfLoVi/DYCi5/x0h/bAOuWF3dULTKpt2Vo= go.yaml.in/yaml/v3 v3.0.4 h1:tfq32ie2Jv2UxXFdLJdh3jXuOzWiL1fo0bu/FbuKpbc= go.yaml.in/yaml/v3 v3.0.4/go.mod h1:DhzuOOF2ATzADvBadXxruRBLzYTpT36CKvDb3+aBEFg= +golang.org/x/exp v0.0.0-20240404231335-c0f41cb1a7a0 h1:985EYyeCOxTpcgOTJpflJUwOeEz0CQOdPt73OzpE9F8= +golang.org/x/exp v0.0.0-20240404231335-c0f41cb1a7a0/go.mod h1:/lliqkxwWAhPjf5oSOIJup2XcqJaw8RGS6k3TGEc7GI= +golang.org/x/net v0.48.0 h1:zyQRTTrjc33Lhh0fBgT/H3oZq9WuvRR5gPC70xpDiQU= +golang.org/x/net v0.48.0/go.mod h1:+ndRgGjkh8FGtu1w1FGbEC31if4VrNVMuKTgcAAnQRY= golang.org/x/oauth2 v0.37.0 h1:JUlcxA8oAtauLfiH8FX2/FkAWHAdi0QtGCGc+hofE98= golang.org/x/oauth2 v0.37.0/go.mod h1:IxwZNxUULJmpBFf9K/9NTMSIfZZuvuTy1gGxhigP/58= -golang.org/x/sys v0.41.0 h1:Ivj+2Cp/ylzLiEU89QhWblYnOE9zerudt9Ftecq2C6k= -golang.org/x/sys v0.41.0/go.mod h1:OgkHotnGiDImocRcuBABYBEXf8A9a87e/uXjp9XT3ks= +golang.org/x/sync v0.22.0 h1:SZjpbeLmrCk4xhRSZFNZW5gFUeCeFgjekvI/+gfScek= +golang.org/x/sync v0.22.0/go.mod h1:9xrNwdLfx4jkKbNva9FpL6vEN7evnE43NNNJQ2LF3+0= +golang.org/x/sys v0.47.0 h1:o7XGOvZQCADBQQ4Y7VNq2dRWQR7JmOUW8Kxx4ZsNgWs= +golang.org/x/sys v0.47.0/go.mod h1:4GL1E5IUh+htKOUEOaiffhrAeqysfVGipDYzABqnCmw= golang.org/x/text v0.34.0 h1:oL/Qq0Kdaqxa1KbNeMKwQq0reLCCaFtqu2eNuSeNHbk= golang.org/x/text v0.34.0/go.mod h1:homfLqTYRFyVYemLBFl5GgL/DWEiH5wcsQ5gSh1yziA= +golang.org/x/time v0.15.0 h1:bbrp8t3bGUeFOx08pvsMYRTCVSMk89u4tKbNOZbp88U= +golang.org/x/time v0.15.0/go.mod h1:Y4YMaQmXwGQZoFaVFk4YpCt4FLQMYKZe9oeV/f4MSno= +golang.org/x/tools v0.42.0 h1:uNgphsn75Tdz5Ji2q36v/nsFSfR/9BRFvqhGBaJGd5k= +golang.org/x/tools v0.42.0/go.mod h1:Ma6lCIwGZvHK6XtgbswSoWroEkhugApmsXyrUmBhfr0= +google.golang.org/protobuf v1.36.11 h1:fV6ZwhNocDyBLK0dj+fg8ektcVegBBuEolpbTQyBNVE= +google.golang.org/protobuf v1.36.11/go.mod h1:HTf+CrKn2C3g5S8VImy6tdcUvCska2kB7j23XfzDpco= gopkg.in/check.v1 v0.0.0-20161208181325-20d25e280405/go.mod h1:Co6ibVJAznAaIkqp8huTwlJQCZ016jof/cbN4VW5Yz0= gopkg.in/check.v1 v1.0.0-20201130134442-10cb98267c6c h1:Hei/4ADfdWqJk1ZMxUNpqntNwaWcugrBjAiHlqqRiVk= gopkg.in/check.v1 v1.0.0-20201130134442-10cb98267c6c/go.mod h1:JHkPIbrfpd72SG/EVd6muEfDQjcINNoR0C8j2r3qZ4Q= gopkg.in/yaml.v3 v3.0.1 h1:fxVm/GzAzEWqLHuvctI91KS9hhNmmWOoWu0XTYJS7CA= gopkg.in/yaml.v3 v3.0.1/go.mod h1:K4uyk7z7BCEPqu6E+C64Yfv1cQ7kz7rIZviUmN+EgEM= +mvdan.cc/sh/v3 v3.14.1 h1:bXkhQWNHCs0KZEChF8hYS6FC+T2N9mUZLbQv9blditI= +mvdan.cc/sh/v3 v3.14.1/go.mod h1:syYCoFET8w9tvevxiXUtY8/ICrU+l26jHmhJDra3Vwo= diff --git a/internal/agent/activity.go b/internal/agent/activity.go new file mode 100644 index 0000000..5f4550c --- /dev/null +++ b/internal/agent/activity.go @@ -0,0 +1,192 @@ +package agent + +import ( + "fmt" + "strings" + "time" + + rootly "github.com/rootlyhq/rootly-go" +) + +type pendingActivity struct { + title string + startedAt time.Time +} + +type completedActivity struct { + line string + finalizedByText bool + finalizedByQualifyingStatus bool +} + +type activityTracker struct { + pending *pendingActivity + completed []completedActivity + taskStarted map[string]time.Time + taskFinished map[string]bool + now func() time.Time +} + +func (tracker *activityTracker) Handle(event Event) []string { + tracker.initialize() + switch event.Type { + case rootly.AiChatStreamEventStatus: + if event.Status == nil { + return nil + } + title, startsActivity := statusActivityTitle(event.Status.Status) + tracker.finishPending(true, false, startsActivity) + if startsActivity { + lines := tracker.flushCompleted() + tracker.pending = &pendingActivity{title: title, startedAt: tracker.now()} + return lines + } + case rootly.AiChatStreamEventText: + tracker.finishPending(true, true, false) + case rootly.AiChatStreamEventDone: + if event.Done != nil && event.Done.Status == "confirmation_required" { + tracker.dropConfirmation() + return nil + } + tracker.finishPending(event.Done != nil && isSuccessfulDoneStatus(event.Done.Status), false, false) + return tracker.flushCompleted() + case rootly.AiChatStreamEventConfirmationRequired: + tracker.dropConfirmation() + case rootly.AiChatStreamEventError: + tracker.finishPending(false, false, false) + return tracker.flushCompleted() + case rootly.AiChatStreamEventTaskUpdate: + return tracker.handleTaskUpdate(event.TaskUpdate) + } + return nil +} + +func isSuccessfulDoneStatus(status string) bool { + return status == "complete" || status == "confirmed" +} + +func (tracker *activityTracker) Finish(success bool) []string { + tracker.initialize() + tracker.finishPending(success, false, false) + return tracker.flushCompleted() +} + +func (tracker *activityTracker) Reset() { + tracker.pending = nil + tracker.completed = nil + tracker.taskStarted = make(map[string]time.Time) + tracker.taskFinished = make(map[string]bool) +} + +func (tracker *activityTracker) discardPendingIf(activity *pendingActivity) { + if activity != nil && tracker.pending == activity { + tracker.pending = nil + } +} + +// completePendingIf marks activity done, as its local tool really ran, and +// returns the lines ready to print. +func (tracker *activityTracker) completePendingIf(activity *pendingActivity) []string { + tracker.initialize() + if activity == nil || tracker.pending != activity { + return nil + } + tracker.finishPending(true, false, false) + return tracker.flushCompleted() +} + +func (tracker *activityTracker) initialize() { + if tracker.now == nil { + tracker.now = time.Now + } + if tracker.taskStarted == nil { + tracker.taskStarted = make(map[string]time.Time) + } + if tracker.taskFinished == nil { + tracker.taskFinished = make(map[string]bool) + } +} + +func (tracker *activityTracker) finishPending(success, finalizedByText, finalizedByQualifyingStatus bool) { + if tracker.pending == nil { + return + } + glyph := "✓" + if !success { + glyph = "✗" + } + line := fmt.Sprintf("%s %s (%.1fs)", glyph, tracker.pending.title, tracker.now().Sub(tracker.pending.startedAt).Seconds()) + tracker.completed = append(tracker.completed, completedActivity{ + line: line, + finalizedByText: finalizedByText, + finalizedByQualifyingStatus: finalizedByQualifyingStatus, + }) + tracker.pending = nil +} + +func (tracker *activityTracker) flushCompleted() []string { + if len(tracker.completed) == 0 { + return nil + } + lines := make([]string, len(tracker.completed)) + for index, activity := range tracker.completed { + lines[index] = activity.line + } + tracker.completed = nil + return lines +} + +func (tracker *activityTracker) dropConfirmation() { + tracker.pending = nil + for len(tracker.completed) > 0 { + last := tracker.completed[len(tracker.completed)-1] + if !last.finalizedByText && last.finalizedByQualifyingStatus { + break + } + tracker.completed = tracker.completed[:len(tracker.completed)-1] + } +} + +func (tracker *activityTracker) handleTaskUpdate(task *rootly.AiChatStreamTaskUpdate) []string { + if task == nil { + return nil + } + if _, exists := tracker.taskStarted[task.ID]; !exists { + tracker.taskStarted[task.ID] = tracker.now() + } + status := strings.ToLower(task.Status) + if status == "in_progress" || status == "pending" || status == "running" || tracker.taskFinished[task.ID] { + return nil + } + tracker.taskFinished[task.ID] = true + glyph := "✓" + if strings.Contains(status, "fail") || strings.Contains(status, "error") || strings.Contains(status, "cancel") { + glyph = "✗" + } + if hasActivityDuration(task.Title) { + return []string{fmt.Sprintf("%s %s", glyph, task.Title)} + } + elapsed := tracker.now().Sub(tracker.taskStarted[task.ID]).Seconds() + return []string{fmt.Sprintf("%s %s (%.1fs)", glyph, task.Title, elapsed)} +} + +func statusActivityTitle(status string) (string, bool) { + status = strings.TrimSpace(status) + if !strings.HasSuffix(status, "...") || strings.EqualFold(status, "Composing response...") { + return "", false + } + title := strings.TrimSpace(strings.TrimSuffix(status, "...")) + return title, title != "" +} + +func hasActivityDuration(title string) bool { + if !strings.HasSuffix(title, ")") { + return false + } + start := strings.LastIndex(title, " (") + if start < 0 { + return false + } + _, err := time.ParseDuration(title[start+2 : len(title)-1]) + return err == nil +} diff --git a/internal/agent/activity_test.go b/internal/agent/activity_test.go new file mode 100644 index 0000000..b70ca20 --- /dev/null +++ b/internal/agent/activity_test.go @@ -0,0 +1,278 @@ +package agent + +import ( + "bytes" + "context" + "strings" + "testing" + "time" + + rootly "github.com/rootlyhq/rootly-go" +) + +func TestActivityTrackerStatusSequence(t *testing.T) { + now := time.Date(2026, 1, 1, 0, 0, 0, 0, time.UTC) + tracker := activityTracker{now: func() time.Time { return now }} + var lines []string + + lines = append(lines, tracker.Handle(Event{ + Type: rootly.AiChatStreamEventStatus, + Status: &rootly.AiChatStreamStatus{Status: "Get incidents..."}, + })...) + now = now.Add(1200 * time.Millisecond) + lines = append(lines, tracker.Handle(Event{ + Type: rootly.AiChatStreamEventStatus, + Status: &rootly.AiChatStreamStatus{Status: "Composing response..."}, + })...) + lines = append(lines, tracker.Handle(Event{Type: rootly.AiChatStreamEventText})...) + lines = append(lines, tracker.Handle(Event{ + Type: rootly.AiChatStreamEventDone, + Done: &rootly.AiChatStreamDone{Status: "complete"}, + })...) + + if len(lines) != 1 || !strings.HasPrefix(lines[0], "✓ Get incidents (") { + t.Fatalf("activity lines = %#v, want one successful Get incidents line", lines) + } +} + +func TestActivityTrackerConfirmedDoneIsSuccessful(t *testing.T) { + tracker := activityTracker{} + tracker.Handle(activityStatusEvent("Get incidents...")) + + lines := tracker.Handle(activityDoneEvent("confirmed")) + if len(lines) != 1 || !strings.HasPrefix(lines[0], "✓ Get incidents (") { + t.Fatalf("activity lines = %#v, want one successful Get incidents line", lines) + } +} + +func TestActivityTrackerCompletesConsecutiveStatusesInOrder(t *testing.T) { + now := time.Date(2026, 1, 1, 0, 0, 0, 0, time.UTC) + tracker := activityTracker{now: func() time.Time { return now }} + var lines []string + + lines = append(lines, tracker.Handle(activityStatusEvent("Get incidents..."))...) + now = now.Add(time.Second) + lines = append(lines, tracker.Handle(activityStatusEvent("Search alerts..."))...) + now = now.Add(time.Second) + lines = append(lines, tracker.Handle(activityDoneEvent("complete"))...) + + if len(lines) != 2 || !strings.HasPrefix(lines[0], "✓ Get incidents (") || !strings.HasPrefix(lines[1], "✓ Search alerts (") { + t.Fatalf("activity lines = %#v, want two successful lines in order", lines) + } +} + +func TestActivityTrackerKeepsRepeatedTitlesAsSeparateActivities(t *testing.T) { + tracker := activityTracker{} + var lines []string + + lines = append(lines, tracker.Handle(activityStatusEvent("Get incidents..."))...) + lines = append(lines, tracker.Handle(activityStatusEvent("Get incidents..."))...) + lines = append(lines, tracker.Handle(activityDoneEvent("complete"))...) + + if len(lines) != 2 || !strings.HasPrefix(lines[0], "✓ Get incidents (") || !strings.HasPrefix(lines[1], "✓ Get incidents (") { + t.Fatalf("activity lines = %#v, want two separate successful activities", lines) + } +} + +func TestActivityTrackerConfirmationDropsTextFinalization(t *testing.T) { + tracker := activityTracker{} + var beforeResume []string + + beforeResume = append(beforeResume, tracker.Handle(activityStatusEvent("A..."))...) + beforeResume = append(beforeResume, tracker.Handle(Event{Type: rootly.AiChatStreamEventText})...) + beforeResume = append(beforeResume, tracker.Handle(Event{Type: rootly.AiChatStreamEventConfirmationRequired})...) + beforeResume = append(beforeResume, tracker.Handle(activityDoneEvent("confirmation_required"))...) + if len(beforeResume) != 0 { + t.Fatalf("activity lines before resume = %#v, want none", beforeResume) + } + + var afterResume []string + afterResume = append(afterResume, tracker.Handle(activityStatusEvent("A..."))...) + afterResume = append(afterResume, tracker.Handle(Event{Type: rootly.AiChatStreamEventText})...) + afterResume = append(afterResume, tracker.Handle(activityDoneEvent("complete"))...) + if len(afterResume) != 1 || !strings.HasPrefix(afterResume[0], "✓ A (") { + t.Fatalf("activity lines after resume = %#v, want one successful A line", afterResume) + } +} + +func TestActivityTrackerFlushesCompletedActivityBeforeConfirmation(t *testing.T) { + tracker := activityTracker{} + tracker.Handle(activityStatusEvent("A...")) + + flushed := tracker.Handle(activityStatusEvent("B...")) + if len(flushed) != 1 || !strings.HasPrefix(flushed[0], "✓ A (") { + t.Fatalf("activity lines on B status = %#v, want one successful A line", flushed) + } + + var afterConfirmation []string + afterConfirmation = append(afterConfirmation, tracker.Handle(Event{Type: rootly.AiChatStreamEventText})...) + afterConfirmation = append(afterConfirmation, tracker.Handle(Event{Type: rootly.AiChatStreamEventConfirmationRequired})...) + afterConfirmation = append(afterConfirmation, tracker.Handle(activityDoneEvent("confirmation_required"))...) + if len(afterConfirmation) != 0 { + t.Fatalf("activity lines for B = %#v, want none", afterConfirmation) + } +} + +func TestActivityTrackerFinishFlushesStatusFinalization(t *testing.T) { + tracker := activityTracker{} + tracker.Handle(activityStatusEvent("A...")) + tracker.Handle(activityStatusEvent("Composing response...")) + + lines := tracker.Finish(true) + if len(lines) != 1 || !strings.HasPrefix(lines[0], "✓ A (") { + t.Fatalf("activity lines on Finish() = %#v, want retained successful A line", lines) + } +} + +func TestActivityTrackerErrorAndConfirmation(t *testing.T) { + t.Run("error", func(t *testing.T) { + tracker := activityTracker{} + tracker.Handle(activityStatusEvent("Get incidents...")) + lines := tracker.Handle(Event{Type: rootly.AiChatStreamEventError, Error: &rootly.AiChatStreamError{Message: "failed"}}) + if len(lines) != 1 || !strings.HasPrefix(lines[0], "✗ Get incidents (") { + t.Fatalf("activity lines = %#v, want failed Get incidents line", lines) + } + }) + + t.Run("confirmation required", func(t *testing.T) { + tracker := activityTracker{} + tracker.Handle(activityStatusEvent("Create action item...")) + lines := tracker.Handle(Event{Type: rootly.AiChatStreamEventConfirmationRequired}) + lines = append(lines, tracker.Handle(activityDoneEvent("confirmation_required"))...) + if len(lines) != 0 { + t.Fatalf("activity lines = %#v, want pending activity dropped", lines) + } + }) +} + +func TestActivityTrackerDoesNotDuplicateTaskDuration(t *testing.T) { + tracker := activityTracker{} + tracker.Handle(Event{ + Type: rootly.AiChatStreamEventTaskUpdate, + TaskUpdate: &rootly.AiChatStreamTaskUpdate{ID: "task-1", Title: "Searched incidents (1.2s)", Status: "in_progress"}, + }) + lines := tracker.Handle(Event{ + Type: rootly.AiChatStreamEventTaskUpdate, + TaskUpdate: &rootly.AiChatStreamTaskUpdate{ID: "task-1", Title: "Searched incidents (1.2s)", Status: "complete"}, + }) + if len(lines) != 1 || lines[0] != "✓ Searched incidents (1.2s)" { + t.Fatalf("activity lines = %#v, want one line without duplicate duration", lines) + } +} + +func TestRunPrintWritesStatusActivityOnlyToTTYTextStderr(t *testing.T) { + for _, test := range []struct { + name string + format string + stderrTTY bool + wantActivity bool + }{ + {name: "TTY text", format: "text", stderrTTY: true, wantActivity: true}, + {name: "piped text", format: "text"}, + {name: "TTY JSON", format: "json", stderrTTY: true}, + } { + t.Run(test.name, func(t *testing.T) { + streamer := &fakeStreamer{responses: [][]rootly.AiChatStreamEvent{{ + eventOf(rootly.AiChatStreamEventStatus, `{"status":"Get incidents..."}`), + eventOf(rootly.AiChatStreamEventStatus, `{"status":"Composing response..."}`), + eventOf(rootly.AiChatStreamEventText, `{"content":"Done"}`), + eventOf(rootly.AiChatStreamEventDone, `{"status":"complete"}`), + }}} + var stdout, stderr strings.Builder + options := PrintOptions{ + Prompt: "list incidents", + Stdin: strings.NewReader(""), + Stdout: &stdout, + Stderr: &stderr, + OutputFormat: test.format, + IsStdinTTY: func() bool { return true }, + IsStderrTTY: func() bool { return test.stderrTTY }, + } + if err := RunPrint(context.Background(), streamer, options); err != nil { + t.Fatalf("RunPrint() error = %v", err) + } + gotActivity := strings.Contains(stderr.String(), "✓ Get incidents (") + if gotActivity != test.wantActivity { + t.Errorf("stderr = %q, want activity output %v", stderr.String(), test.wantActivity) + } + if strings.Contains(stdout.String(), "✓ Get incidents") { + t.Errorf("stdout = %q, activity must never be printed to stdout", stdout.String()) + } + }) + } +} + +func TestRunPrintDefersConfirmationActivityUntilApproval(t *testing.T) { + streamer := &fakeStreamer{responses: [][]rootly.AiChatStreamEvent{ + { + eventOf(rootly.AiChatStreamEventSessionID, `{"session_id":"`+testSessionID+`"}`), + eventOf(rootly.AiChatStreamEventStatus, `{"status":"Create action item..."}`), + eventOf(rootly.AiChatStreamEventStatus, `{"status":"Composing response..."}`), + eventOf(rootly.AiChatStreamEventText, `{"content":"Ready for approval."}`), + eventOf(rootly.AiChatStreamEventConfirmationRequired, `{"checkpoint_id":"`+testCheckpointID+`","question":"Create the action item?","preview":{"title":"Create action item"},"expires_at":"2099-01-01T00:00:00Z"}`), + eventOf(rootly.AiChatStreamEventDone, `{"status":"confirmation_required"}`), + }, + { + eventOf(rootly.AiChatStreamEventStatus, `{"status":"Create action item..."}`), + eventOf(rootly.AiChatStreamEventText, `{"content":"Created."}`), + eventOf(rootly.AiChatStreamEventDone, `{"status":"complete"}`), + }, + }} + var stdout, stderr bytes.Buffer + options := newPrintOptions(&stdout, &stderr) + options.Stdin = strings.NewReader("y\n") + options.IsStdinTTY = func() bool { return true } + options.IsStderrTTY = func() bool { return true } + + if err := RunPrint(context.Background(), streamer, options); err != nil { + t.Fatalf("RunPrint() error = %v", err) + } + activity := "✓ Create action item (" + if got := strings.Count(stderr.String(), activity); got != 1 { + t.Fatalf("activity output count = %d, want one: %q", got, stderr.String()) + } + if strings.Index(stderr.String(), activity) < strings.Index(stderr.String(), "Approve? [y/N]") { + t.Errorf("activity was printed before approval: %q", stderr.String()) + } +} + +func TestRunPrintFlushesStatusFinalizationAtConfirmation(t *testing.T) { + streamer := &fakeStreamer{responses: [][]rootly.AiChatStreamEvent{ + { + eventOf(rootly.AiChatStreamEventSessionID, `{"session_id":"`+testSessionID+`"}`), + eventOf(rootly.AiChatStreamEventStatus, `{"status":"Create action item..."}`), + eventOf(rootly.AiChatStreamEventStatus, `{"status":"Save incident..."}`), + eventOf(rootly.AiChatStreamEventConfirmationRequired, `{"checkpoint_id":"`+testCheckpointID+`","question":"Create the action item?","preview":{"title":"Create action item"},"expires_at":"2099-01-01T00:00:00Z"}`), + eventOf(rootly.AiChatStreamEventDone, `{"status":"confirmation_required"}`), + }, + { + eventOf(rootly.AiChatStreamEventDone, `{"status":"complete"}`), + }, + }} + var stdout, stderr bytes.Buffer + options := newPrintOptions(&stdout, &stderr) + options.Stdin = strings.NewReader("y\n") + options.IsStdinTTY = func() bool { return true } + options.IsStderrTTY = func() bool { return true } + + if err := RunPrint(context.Background(), streamer, options); err != nil { + t.Fatalf("RunPrint() error = %v", err) + } + activity := "✓ Create action item (" + activityIndex := strings.Index(stderr.String(), activity) + promptIndex := strings.Index(stderr.String(), "Approve? [y/N]") + if activityIndex < 0 || promptIndex < 0 || activityIndex > promptIndex { + t.Errorf("stderr = %q, want completed tool activity before the approval prompt", stderr.String()) + } + if strings.Contains(stderr.String(), "✓ Save incident (") { + t.Errorf("stderr = %q, want pending action activity dropped", stderr.String()) + } +} + +func activityStatusEvent(status string) Event { + return Event{Type: rootly.AiChatStreamEventStatus, Status: &rootly.AiChatStreamStatus{Status: status}} +} + +func activityDoneEvent(status string) Event { + return Event{Type: rootly.AiChatStreamEventDone, Done: &rootly.AiChatStreamDone{Status: status}} +} diff --git a/internal/agent/confirm_card.go b/internal/agent/confirm_card.go new file mode 100644 index 0000000..db74f64 --- /dev/null +++ b/internal/agent/confirm_card.go @@ -0,0 +1,223 @@ +package agent + +import ( + "bytes" + "encoding/json" + "fmt" + "strings" + "time" + + "charm.land/lipgloss/v2" + rootly "github.com/rootlyhq/rootly-go" +) + +type confirmationPreview struct { + Title string `json:"title"` + Summary string `json:"summary"` + ContextText string `json:"context_text"` + Fields []struct { + Label string `json:"label"` + Value any `json:"value"` + Inferred bool `json:"inferred"` + } `json:"fields"` + Destructive bool `json:"destructive"` + InferredFields json.RawMessage `json:"inferred_fields"` +} + +const confirmationKeys = "y approve · n/Esc cancel" + +func renderConfirmationCard(confirmation *rootly.AiChatStreamConfirmationRequired, width int, noColor bool, now time.Time) string { + return renderConfirmationCardWithKeys(confirmation, width, noColor, now, confirmationKeys) +} + +func renderConfirmationCardWithKeys(confirmation *rootly.AiChatStreamConfirmationRequired, width int, noColor bool, now time.Time, keys string) string { + if confirmation == nil { + return "" + } + + lines := []string{sanitizeTerminalText(confirmation.Question)} + preview, validPreview := parseConfirmationPreview(confirmation.Preview) + if validPreview { + if preview.Title != "" { + lines = append(lines, sanitizeTerminalText(preview.Title)) + } + if preview.Summary != "" { + lines = append(lines, sanitizeTerminalText(preview.Summary)) + } + if preview.ContextText != "" { + lines = append(lines, sanitizeTerminalText(preview.ContextText)) + } + inferredFields := confirmationInferredFields(preview.InferredFields) + for _, field := range preview.Fields { + label := sanitizeTerminalText(field.Label) + value := sanitizeTerminalText(confirmationValue(field.Value)) + row := fmt.Sprintf("%s: %s", label, value) + if field.Inferred || inferredFields[label] { + if noColor { + row += " (inferred)" + } else { + row += lipgloss.NewStyle().Faint(true).Render(" (inferred)") + } + } + lines = append(lines, row) + } + if preview.Destructive { + tag := "DESTRUCTIVE" + if !noColor { + tag = lipgloss.NewStyle().Foreground(lipgloss.Color("1")).Bold(true).Render(tag) + } + lines = append(lines, tag) + } + } else if len(confirmation.Preview) > 0 { + lines = append(lines, strings.Split(sanitizedConfirmationPreview(confirmation.Preview), "\n")...) + } + missing := confirmationMissingFields(confirmation.MissingFields) + for index := range missing { + missing[index] = sanitizeTerminalText(missing[index]) + } + if len(missing) > 0 { + lines = append(lines, "Missing: "+strings.Join(missing, ", ")) + } + + footer := keys + if remaining, ok := confirmationRemaining(confirmation.ExpiresAt, now); ok { + footer += " · expires in " + remaining + } + if !noColor { + lines[0] = lipgloss.NewStyle().Bold(true).Render(lines[0]) + footer = lipgloss.NewStyle().Faint(true).Render(footer) + } + lines = append(lines, footer) + + content := strings.Join(lines, "\n") + if noColor { + return indent(content, " ") + } + + cardWidth := min(max(width, 1), 80) + innerWidth := max(1, cardWidth-4) + return lipgloss.NewStyle(). + Border(lipgloss.RoundedBorder()). + Padding(0, 1). + Width(innerWidth). + MaxWidth(innerWidth). + Render(content) +} + +func parseConfirmationPreview(raw json.RawMessage) (confirmationPreview, bool) { + var preview confirmationPreview + if len(raw) == 0 || json.Unmarshal(raw, &preview) != nil || len(preview.Fields) == 0 { + return confirmationPreview{}, false + } + return preview, true +} + +func confirmationValue(value any) string { + switch value := value.(type) { + case nil: + return "" + case string: + return value + case map[string]any, []any: + encoded, err := json.Marshal(sanitizeConfirmationJSONValue(value)) + if err == nil { + return string(encoded) + } + } + return fmt.Sprint(value) +} + +func sanitizedConfirmationPreview(raw []byte) string { + var value any + if json.Unmarshal(raw, &value) != nil { + return sanitizeTerminalText(string(raw)) + } + encoded, err := json.Marshal(sanitizeConfirmationJSONValue(value)) + if err != nil { + return sanitizeTerminalText(string(raw)) + } + var indented bytes.Buffer + if err := json.Indent(&indented, encoded, "", " "); err != nil { + return sanitizeTerminalText(string(encoded)) + } + return indented.String() +} + +func sanitizeConfirmationJSONValue(value any) any { + switch value := value.(type) { + case string: + return sanitizeTerminalText(value) + case []any: + sanitized := make([]any, len(value)) + for index, item := range value { + sanitized[index] = sanitizeConfirmationJSONValue(item) + } + return sanitized + case map[string]any: + sanitized := make(map[string]any, len(value)) + for key, item := range value { + sanitized[sanitizeTerminalText(key)] = sanitizeConfirmationJSONValue(item) + } + return sanitized + default: + return value + } +} + +func confirmationInferredFields(raw json.RawMessage) map[string]bool { + inferred := make(map[string]bool) + var values []json.RawMessage + if json.Unmarshal(raw, &values) != nil { + return inferred + } + for _, value := range values { + var name string + if json.Unmarshal(value, &name) == nil { + inferred[sanitizeTerminalText(name)] = true + continue + } + var field map[string]any + if json.Unmarshal(value, &field) == nil { + if name, ok := field["label"].(string); ok { + inferred[sanitizeTerminalText(name)] = true + } else if name, ok := field["name"].(string); ok { + inferred[sanitizeTerminalText(name)] = true + } + } + } + return inferred +} + +func confirmationMissingFields(raw json.RawMessage) []string { + var fields []map[string]any + if json.Unmarshal(raw, &fields) != nil { + return nil + } + missing := make([]string, 0, len(fields)) + for _, field := range fields { + if label, ok := field["label"].(string); ok && label != "" { + missing = append(missing, label) + } else if name, ok := field["name"].(string); ok && name != "" { + missing = append(missing, name) + } + } + return missing +} + +func confirmationRemaining(expiresAt string, now time.Time) (string, bool) { + if expiresAt == "" { + return "", false + } + expiration, err := time.Parse(time.RFC3339, expiresAt) + if err != nil { + return "", false + } + remaining := expiration.Sub(now) + if remaining >= time.Minute { + return fmt.Sprintf("%dm", int((remaining+time.Minute-1)/time.Minute)), true + } + if remaining <= 0 { + return "0s", true + } + return fmt.Sprintf("%ds", int((remaining+time.Second-1)/time.Second)), true +} diff --git a/internal/agent/confirm_card_test.go b/internal/agent/confirm_card_test.go new file mode 100644 index 0000000..e187084 --- /dev/null +++ b/internal/agent/confirm_card_test.go @@ -0,0 +1,126 @@ +package agent + +import ( + "strings" + "testing" + "time" + + "charm.land/lipgloss/v2" + "github.com/charmbracelet/x/ansi" + rootly "github.com/rootlyhq/rootly-go" +) + +func TestRenderConfirmationCard(t *testing.T) { + now := time.Date(2026, time.January, 2, 3, 4, 5, 0, time.UTC) + confirmation := &rootly.AiChatStreamConfirmationRequired{ + Question: "Create the incident?", + Preview: []byte(`{ + "title":"Incident", + "summary":"Create an incident from this alert.", + "context_text":"Triggered by the on-call team.", + "fields":[ + {"label":"Title","value":"Database unavailable"}, + {"label":"Severity","value":"P1","inferred":true} + ], + "destructive":true, + "inferred_fields":["Owner"] + }`), + MissingFields: []byte(`[{"label":"Service"},{"name":"Environment"}]`), + ExpiresAt: now.Add(4*time.Minute + 2*time.Second).Format(time.RFC3339), + } + + colored := renderConfirmationCard(confirmation, 40, false, now) + flattened := strings.Join(strings.Fields(ansi.Strip(colored)), " ") + for _, expected := range []string{ + "Create the incident?", + "Incident", + "Create an incident from this", + "alert.", + "Triggered by the on-call team.", + "Title: Database unavailable", + "Severity: P1", + "(inferred)", + "Missing: Service, Environment", + "DESTRUCTIVE", + "expires in 5m", + "╭", + } { + if !strings.Contains(flattened, expected) { + t.Errorf("colored card = %q, want it to contain %q", colored, expected) + } + } + for _, line := range strings.Split(colored, "\n") { + if width := lipgloss.Width(line); width > 40 { + t.Errorf("colored card line width = %d, want at most 40: %q", width, line) + } + } + + plain := renderConfirmationCard(confirmation, 40, true, now) + if strings.Contains(plain, "\x1b[") || strings.Contains(plain, "╭") { + t.Errorf("no-color card contains terminal styling: %q", plain) + } + if !strings.HasPrefix(plain, " Create the incident?") || !strings.Contains(plain, " Severity: P1 (inferred)") { + t.Errorf("no-color card = %q, want indented plain text", plain) + } +} + +func TestRenderConfirmationCardFallsBackToIndentedJSON(t *testing.T) { + for _, preview := range [][]byte{ + []byte(`{"name":"api"}`), + []byte(`not-json`), + } { + confirmation := &rootly.AiChatStreamConfirmationRequired{ + Question: "Approve?", + Preview: preview, + } + got := renderConfirmationCard(confirmation, 80, true, time.Now()) + if !strings.HasPrefix(got, " Approve?\n ") { + t.Errorf("card = %q, want question and indented fallback", got) + } + } +} + +func TestRenderConfirmationCardRendersMissingFieldsWithoutPreviewFields(t *testing.T) { + confirmation := &rootly.AiChatStreamConfirmationRequired{ + Question: "Approve?", + Preview: []byte(`{"title":"Create"}`), + MissingFields: []byte(`[{"label":"Service \u001b[2J"}]`), + } + + got := renderConfirmationCard(confirmation, 80, true, time.Now()) + if !strings.Contains(got, "Missing: Service") { + t.Errorf("card = %q, want missing fields in the fallback branch", got) + } + if strings.Contains(got, "\x1b") { + t.Errorf("card = %q, want sanitized missing fields", got) + } +} + +func TestRenderConfirmationCardOmitsExpiryWhenUnset(t *testing.T) { + got := renderConfirmationCard(&rootly.AiChatStreamConfirmationRequired{ + Question: "Approve?", + }, 80, true, time.Now()) + if strings.Contains(got, "expires in") { + t.Errorf("card = %q, want no expiry countdown", got) + } +} + +func TestRenderConfirmationCardSanitizesTerminalControls(t *testing.T) { + const attack = "\x1b]52;c;ZXZpbA==\x07\x1b[2J\x1b]8;;https://evil.example\x07" + confirmation := &rootly.AiChatStreamConfirmationRequired{ + Question: "Approve " + attack, + Preview: []byte(`{ + "title":"Title \u001b]52;c;ZXZpbA==\u0007\u001b[2J", + "summary":"Summary \u001b]52;c;ZXZpbA==\u0007", + "context_text":"Context \u001b[2J", + "fields":[{"label":"Label \u001b]52;c;ZXZpbA==\u0007","value":{"key":"Value \u001b]52;c;ZXZpbA==\u0007"}}] + }`), + MissingFields: []byte(`["Missing \u001b]52;c;ZXZpbA==\u0007"]`), + } + card := renderConfirmationCard(confirmation, 80, true, time.Now()) + for _, unwanted := range []string{"]52;", "[2J", "\x1b]8;;", "\x1b"} { + if strings.Contains(card, unwanted) { + t.Errorf("confirmation card = %q, want no %q", card, unwanted) + } + } +} diff --git a/internal/agent/credhint.go b/internal/agent/credhint.go new file mode 100644 index 0000000..01e9dc7 --- /dev/null +++ b/internal/agent/credhint.go @@ -0,0 +1,25 @@ +package agent + +import ( + "regexp" + + "github.com/rootlyhq/rootly-cli/internal/mcp" +) + +const localCredsHint = "Refresh local creds, e.g. `aws sso login`" + +var ( + expiredLocalCredsPattern = regexp.MustCompile(`(?i)(sso session .*(expired|invalid)|token (has|is) expired|expiredtoken|expired ?token|credentials? (have|has) expired|aws sso login|unable to locate credentials)`) + deniedCredsPattern = regexp.MustCompile(`(?i)(\b401\b|\b403\b|unauthori[sz]ed|forbidden|invalid[_ ]token|access denied)`) +) + +// looksLikeExpiredCreds reports whether a local tool result resembles stale local credentials. +func looksLikeExpiredCreds(result mcp.CallResult) bool { + if result.Output == userDeclinedOutput { + return false + } + if expiredLocalCredsPattern.MatchString(result.Output) { + return true + } + return result.IsError && deniedCredsPattern.MatchString(result.Output) +} diff --git a/internal/agent/credhint_test.go b/internal/agent/credhint_test.go new file mode 100644 index 0000000..e777c68 --- /dev/null +++ b/internal/agent/credhint_test.go @@ -0,0 +1,56 @@ +package agent + +import ( + "context" + "strings" + "testing" + + tea "charm.land/bubbletea/v2" + + "github.com/rootlyhq/rootly-cli/internal/mcp" +) + +func TestLooksLikeExpiredCreds(t *testing.T) { + cases := []struct { + result mcp.CallResult + want bool + }{ + {mcp.CallResult{Output: "Error when retrieving token from sso: Token has expired and refresh failed"}, true}, + {mcp.CallResult{Output: "The SSO session associated with this profile has expired or is otherwise invalid."}, true}, + {mcp.CallResult{Output: "An error occurred (ExpiredToken) when calling the GetCallerIdentity operation"}, true}, + {mcp.CallResult{Output: "HTTP 401 Unauthorized", IsError: true}, true}, + {mcp.CallResult{Output: "request failed: 403 Forbidden", IsError: true}, true}, + {mcp.CallResult{Output: "found 401 rows in table forbidden_words"}, false}, + {mcp.CallResult{Output: userDeclinedOutput, IsError: true}, false}, + {mcp.CallResult{Output: "file not found", IsError: true}, false}, + } + for _, tc := range cases { + if got := looksLikeExpiredCreds(tc.result); got != tc.want { + t.Errorf("looksLikeExpiredCreds(%q, error=%v) = %v, want %v", tc.result.Output, tc.result.IsError, got, tc.want) + } + } +} + +func TestCredHintShownOnceInTUI(t *testing.T) { + model := NewTUIModel(context.Background(), &fakeStreamer{}, TUIOptions{NoColor: true}) + model.Update(tea.WindowSizeMsg{Width: 100, Height: 30}) + model.Init() + model.Update(localToolDoneMsg{turn: model.turn, credHint: true}) + model.Update(localToolDoneMsg{turn: model.turn, credHint: true}) + model.Update(localToolDoneMsg{turn: model.turn}) + if got := strings.Count(model.Transcript(), localCredsHint); got != 1 { + t.Fatalf("hint shown %d times, want once:\n%s", got, model.Transcript()) + } +} + +func TestCredHintShownOnceInPrintMode(t *testing.T) { + tools := &fakeLocalTools{ + tools: map[string]mcp.ToolRef{"local_aws_whoami": {ServerID: "user:aws", ServerName: "aws", Tool: "whoami", ReadOnly: true}}, + output: "The SSO session associated with this profile has expired or is otherwise invalid.", + } + trip := newToolRoundTrip(t, 2, "local_aws_whoami") + _, stderr, _ := trip.run(t, tools, true, "") + if got := strings.Count(stderr, localCredsHint); got != 1 { + t.Fatalf("hint shown %d times, want once:\n%s", got, stderr) + } +} diff --git a/internal/agent/exit.go b/internal/agent/exit.go new file mode 100644 index 0000000..6d40ad2 --- /dev/null +++ b/internal/agent/exit.go @@ -0,0 +1,42 @@ +package agent + +import "errors" + +// ExitNotAuthenticated is the print-mode exit code for missing or rejected credentials. +const ExitNotAuthenticated = 2 + +type notAuthenticatedError struct{ message string } + +func (e notAuthenticatedError) Error() string { return e.message } + +// NotAuthenticatedError reports missing or rejected Rootly credentials. +func NotAuthenticatedError(message string) error { + return notAuthenticatedError{message: message} +} + +// IsNotAuthenticated reports whether err comes from missing or rejected credentials. +func IsNotAuthenticated(err error) bool { + var target notAuthenticatedError + return errors.As(err, &target) +} + +// ExitCodeError carries a process exit code for the command runner. +type ExitCodeError struct { + Err error + Code int +} + +func (e ExitCodeError) Error() string { return e.Err.Error() } + +func (e ExitCodeError) Unwrap() error { return e.Err } + +// ExitCode returns the process exit code. +func (e ExitCodeError) ExitCode() int { return e.Code } + +// PrintExitError maps print-mode errors to exit codes: 2 when not authenticated, 1 otherwise. +func PrintExitError(err error) error { + if err != nil && IsNotAuthenticated(err) { + return ExitCodeError{Err: err, Code: ExitNotAuthenticated} + } + return err +} diff --git a/internal/agent/link_test.go b/internal/agent/link_test.go new file mode 100644 index 0000000..3d7b82d --- /dev/null +++ b/internal/agent/link_test.go @@ -0,0 +1,278 @@ +package agent + +import ( + "bytes" + "strings" + "testing" + + "charm.land/glamour/v2" + "charm.land/lipgloss/v2" + "github.com/charmbracelet/x/ansi" +) + +func TestGlamourRendersLinksWithPrivateUseMarkers(t *testing.T) { + const target = "http://localhost:3000/account/incidents/8-iam-credentials-has-been-leaked?tab=follow_ups&ai_slack_unfurl_resource=action_item:e052c503-dc3b-447e-8a5e-e4dd4d2a266a" + markdown := "Created: [Rotate the keys](" + target + ") on INC-8.\n\n- [Review the action](https://rootly.com/docs/actions)" + renderer, err := glamour.NewTermRenderer(glamour.WithEnvironmentConfig(), glamour.WithWordWrap(80)) + if err != nil { + t.Fatal(err) + } + + source, urls := preprocessMarkdownLinksWithURLs(markdown, "api.rootly.com") + rendered, err := renderer.Render(source) + if err != nil { + t.Fatal(err) + } + output := restoreMarkdownLinks(rendered, urls) + visible := ansi.Strip(output) + if !strings.Contains(visible, "Rotate the keys") { + t.Errorf("visible output = %q, want link label", visible) + } + for _, unwanted := range []string{"follow_ups", "ai_slack_unfurl", "\uE000", "\uE001"} { + if strings.Contains(visible, unwanted) { + t.Errorf("visible output = %q, want no %q", visible, unwanted) + } + } + if got := strings.Count(output, ansi.SetHyperlink(target)); got != 1 { + t.Errorf("hyperlink start count = %d, want exactly one", got) + } +} + +func TestRestoreMarkdownLinksStylesHyperlinkText(t *testing.T) { + const destination = "https://rootly.com/account?chat=x" + rendered := "\uE000\x1b[1mopen in web\x1b[0m\uE001" + output := restoreMarkdownLinks(rendered, []markdownLink{{ + destination: destination, + hyperlink: true, + }}) + + start := ansi.SetHyperlink(destination) + end := ansi.ResetHyperlink() + linkStart := strings.Index(output, start) + if linkStart < 0 { + t.Fatalf("output = %q, want OSC-8 open sequence", output) + } + linkEnd := strings.Index(output[linkStart+len(start):], end) + if linkEnd < 0 { + t.Fatalf("output = %q, want OSC-8 close sequence", output) + } + styledText := output[linkStart+len(start) : linkStart+len(start)+linkEnd] + if !strings.Contains(styledText, "\x1b[4") { + t.Errorf("link text = %q, want underline SGR", styledText) + } + wantStyledText := lipgloss.NewStyle(). + Underline(true). + Foreground(lipgloss.Color(metricChartColor)). + Render("open in web") + if !strings.Contains(styledText, wantStyledText) { + t.Errorf("link text = %q, want foreground %s and underlined label", styledText, metricChartColor) + } + if got := strings.Count(ansi.Strip(output), "open in web"); got != 1 { + t.Errorf("visible link label count = %d, want 1: %q", got, ansi.Strip(output)) + } +} + +func TestRestoreMarkdownLinksDoesNotUnderlineWrappedIndentation(t *testing.T) { + rendered := "\uE000first line\n \x1b[1msecond line\x1b[0m\uE001" + output := restoreMarkdownLinks(rendered, []markdownLink{{ + destination: "https://rootly.com/wrapped", + hyperlink: true, + }}) + lines := strings.Split(output, "\n") + if len(lines) != 2 { + t.Fatalf("output lines = %d, want 2: %q", len(lines), output) + } + if !strings.HasPrefix(lines[1], " ") { + t.Fatalf("wrapped line = %q, want leading indentation before style", lines[1]) + } + if !strings.HasPrefix(strings.TrimLeft(lines[1], " "), "\x1b[") { + t.Errorf("wrapped line = %q, want style after its indentation", lines[1]) + } + if got := ansi.Strip(output); got != "first line\n second line" { + t.Errorf("visible output = %q, want wrapped label with indentation", got) + } +} + +func TestGlamourMapsMultipleLinksToURLsInOrder(t *testing.T) { + const firstURL = "https://rootly.com/incidents/first" + const secondURL = "https://rootly.com/incidents/second" + renderer, err := glamour.NewTermRenderer(glamour.WithEnvironmentConfig(), glamour.WithWordWrap(80)) + if err != nil { + t.Fatal(err) + } + markdown := "[First](https://rootly.com/incidents/first) and [Second](https://rootly.com/incidents/second)." + source, urls := preprocessMarkdownLinksWithURLs(markdown, "api.rootly.com") + rendered, err := renderer.Render(source) + if err != nil { + t.Fatal(err) + } + output := restoreMarkdownLinks(rendered, urls) + firstIndex := strings.Index(output, ansi.SetHyperlink(firstURL)) + secondIndex := strings.Index(output, ansi.SetHyperlink(secondURL)) + if firstIndex < 0 || secondIndex <= firstIndex || + strings.Count(output, ansi.SetHyperlink(firstURL)) != 1 || + strings.Count(output, ansi.SetHyperlink(secondURL)) != 1 { + t.Errorf("hyperlinks are missing or reordered: %q", output) + } +} + +func TestGlamourDisclosesUntrustedLinkHostsAndTrustsRootlyDomain(t *testing.T) { + markdown := "[Open INC-8](https://evil.example/x) and [Open Rootly](https://rootly.com/incidents/8)" + renderer, err := glamour.NewTermRenderer(glamour.WithEnvironmentConfig(), glamour.WithWordWrap(80)) + if err != nil { + t.Fatal(err) + } + source, links := preprocessMarkdownLinksWithURLs(markdown, "api.rootly.com") + rendered, err := renderer.Render(source) + if err != nil { + t.Fatal(err) + } + output := restoreMarkdownLinks(rendered, links) + visible := ansi.Strip(output) + if !strings.Contains(visible, "Open INC-8 (evil.example)") { + t.Errorf("visible output = %q, want external host after link label", visible) + } + if !strings.Contains(visible, "Open Rootly") || strings.Contains(visible, "(rootly.com)") { + t.Errorf("visible output = %q, want trusted Rootly link without host disclosure", visible) + } +} + +func TestGlamourDoesNotHyperlinkUnsupportedSchemes(t *testing.T) { + markdown := "[x](javascript:alert(1))" + renderer, err := glamour.NewTermRenderer(glamour.WithEnvironmentConfig(), glamour.WithWordWrap(80)) + if err != nil { + t.Fatal(err) + } + source, links := preprocessMarkdownLinksWithURLs(markdown, "api.rootly.com") + rendered, err := renderer.Render(source) + if err != nil { + t.Fatal(err) + } + output := restoreMarkdownLinks(rendered, links) + if strings.Contains(output, "\x1b]8;;") { + t.Errorf("output = %q, want no OSC-8 hyperlink", output) + } + if visible := ansi.Strip(output); !strings.Contains(visible, "x ") { + t.Errorf("visible output = %q, want label followed by plain URL", visible) + } +} + +func TestGlamourPreservesPrivateUseLinkMarkers(t *testing.T) { + const ( + start = "\uE000" + end = "\uE001" + ) + source := "Rotate " + start + "the keys" + end + " today.\n\n- " + start + "review access" + end + renderer, err := glamour.NewTermRenderer(glamour.WithEnvironmentConfig(), glamour.WithWordWrap(80)) + if err != nil { + t.Fatal(err) + } + rendered, err := renderer.Render(source) + if err != nil { + t.Fatal(err) + } + if strings.Count(rendered, start) != 2 || strings.Count(rendered, end) != 2 { + t.Fatalf("rendered output = %q, want both pairs of private-use markers", rendered) + } + firstStart := strings.Index(rendered, start) + firstEnd := strings.Index(rendered, end) + secondStart := strings.LastIndex(rendered, start) + secondEnd := strings.LastIndex(rendered, end) + if firstStart >= firstEnd || firstEnd >= secondStart || secondStart >= secondEnd { + t.Errorf("rendered marker order is incorrect: %q", rendered) + } +} + +func TestPreprocessMarkdownLinksUsesASTAndHonorsNoColor(t *testing.T) { + const ( + url = "https://rootly.com/docs" + code = "`[code](https://rootly.com/code)`" + block = "[block](https://rootly.com/block)" + autolink = "" + bare = "https://rootly.com/bare" + ) + markdown := "[Read the docs](" + url + ") " + code + "\n\n```md\n" + block + "\n```\n\n" + autolink + " " + bare + + processed := terminalMarkdown(markdown, false) + if !strings.Contains(processed, "\uE000Read the docs\uE001") { + t.Errorf("processed Markdown = %q, want private-use markers around the link text", processed) + } + if strings.Contains(processed, ansi.SetHyperlink(url)) || strings.Contains(processed, ansi.ResetHyperlink()) { + t.Errorf("processed Markdown = %q, must not contain OSC-8 controls before rendering", processed) + } + for _, unchanged := range []string{code, block, autolink, bare} { + if !strings.Contains(processed, unchanged) { + t.Errorf("processed Markdown = %q, want unchanged %q", processed, unchanged) + } + } + if got := terminalMarkdown(markdown, true); got != markdown { + t.Errorf("no-color Markdown = %q, want raw Markdown %q", got, markdown) + } +} + +func TestRenderMarkdownStylesRewrittenDashboardLink(t *testing.T) { + model := NewTUIModel(t.Context(), &fakeStreamer{}, TUIOptions{APIHost: "api.rootly.com"}) + output := model.renderMarkdown("[Add to dashboard](#ai-add-to-dashboard:ref)", 80, false) + if !strings.Contains(output, ansi.SetHyperlink("https://rootly.com/account")) { + t.Errorf("rendered link = %q, want rewritten OSC-8 destination", output) + } + if !strings.Contains(output, "\x1b[4") { + t.Errorf("rendered link = %q, want underlined clickable text", output) + } + if visible := ansi.Strip(output); !strings.Contains(visible, "Add to dashboard: open in web") { + t.Errorf("rendered link = %q, want rewritten label", visible) + } +} + +func TestRenderMarkdownStylesOpenInWebLink(t *testing.T) { + const destination = "https://rootly.com/account?chat=x" + model := NewTUIModel(t.Context(), &fakeStreamer{}, TUIOptions{APIHost: "api.rootly.com"}) + output := model.renderMarkdown("[open in web]("+destination+")", 80, false) + if !strings.Contains(output, ansi.SetHyperlink(destination)) || + !strings.Contains(output, ansi.ResetHyperlink()) { + t.Errorf("rendered link = %q, want OSC-8 open and close sequences", output) + } + if !strings.Contains(output, "\x1b[4") { + t.Errorf("rendered link = %q, want underlined clickable text", output) + } + if got := strings.Count(ansi.Strip(output), "open in web"); got != 1 { + t.Errorf("visible link label count = %d, want 1: %q", got, ansi.Strip(output)) + } +} + +func TestNoColorAndPipedMarkdownSkipLinkStyling(t *testing.T) { + const markdown = "[Add to dashboard](#ai-add-to-dashboard:ref)" + model := NewTUIModel(t.Context(), &fakeStreamer{}, TUIOptions{ + APIHost: "api.rootly.com", + NoColor: true, + }) + tuiOutput := model.renderMarkdown(markdown, 80, false) + if strings.Contains(tuiOutput, "\x1b") { + t.Errorf("no-color TUI output = %q, want no ANSI", tuiOutput) + } + + for _, test := range []struct { + name string + noColor bool + stdoutIsTTY bool + }{ + {name: "no-color terminal", noColor: true, stdoutIsTTY: true}, + {name: "piped output", stdoutIsTTY: false}, + } { + t.Run(test.name, func(t *testing.T) { + var stdout bytes.Buffer + err := writeAnswer(PrintOptions{ + Stdout: &stdout, + APIHost: "api.rootly.com", + NoColor: test.noColor, + IsStdoutTTY: func() bool { return test.stdoutIsTTY }, + }, markdown) + if err != nil { + t.Fatal(err) + } + if strings.Contains(stdout.String(), "\x1b") { + t.Errorf("print output = %q, want no ANSI", stdout.String()) + } + }) + } +} diff --git a/internal/agent/localtools.go b/internal/agent/localtools.go new file mode 100644 index 0000000..0b160c4 --- /dev/null +++ b/internal/agent/localtools.go @@ -0,0 +1,155 @@ +package agent + +import ( + "context" + "encoding/json" + "strings" + "time" + + rootly "github.com/rootlyhq/rootly-go" + + "github.com/rootlyhq/rootly-cli/internal/api" + "github.com/rootlyhq/rootly-cli/internal/mcp" +) + +const ( + localToolsOffMessage = "Local MCP tools are off for your org" + userDeclinedOutput = "User declined" +) + +// LocalTools is the agent's view of the local MCP manager. +type LocalTools interface { + ClientTools() (tools []mcp.ClientTool, omitted int) + Lookup(string) (mcp.ToolRef, bool) + NeedsApproval(string, map[string]any) bool + Call(context.Context, string, map[string]any) mcp.CallResult +} + +// LocalApprovalMemory is implemented by managers that can remember approvals. +type LocalApprovalMemory interface { + ApprovalOptions(string, map[string]any) mcp.ApprovalOptions + RememberApproval(string, map[string]any, bool) error +} + +// LocalReadOnlyCalls is implemented by managers that can tell a read from a +// write allowed by a remembered approval. +type LocalReadOnlyCalls interface { + ReadOnlyCall(string, map[string]any) bool +} + +type ToolResultPoster interface { + PostAiChatToolResult(context.Context, api.ToolResult) error +} + +func localToolLabel(tools LocalTools, call *ClientToolCall) string { + if ref, ok := resolveLocalCall(tools, call); ok { + return "local · " + ref.ServerName + " · " + ref.Tool + } + return "local · " + call.Name +} + +func localToolFailure(tools LocalTools, call *ClientToolCall, output string) string { + label := localToolLabel(tools, call) + reason, _, _ := strings.Cut(strings.TrimSpace(output), "\n") + if runes := []rune(reason); len(runes) > 160 { + reason = string(runes[:160]) + "…" + } + if reason == "" { + return "✗ " + label + " failed" + } + return "✗ " + label + " failed: " + reason +} + +func localToolActivity(ref mcp.ToolRef) string { + return "⚙ local · " + sanitizeTerminalText(ref.ServerName) + " · " + sanitizeTerminalText(ref.Tool) +} + +func resolveLocalCall(tools LocalTools, call *ClientToolCall) (mcp.ToolRef, bool) { + if tools == nil || call == nil { + return mcp.ToolRef{}, false + } + return tools.Lookup(call.Name) +} + +func runLocalCall(ctx context.Context, tools LocalTools, call *ClientToolCall) mcp.CallResult { + if _, ok := resolveLocalCall(tools, call); !ok { + return mcp.CallResult{Output: "Unknown local tool " + call.Name, IsError: true} + } + return tools.Call(ctx, call.Name, call.Arguments) +} + +func postLocalResult(ctx context.Context, poster ToolResultPoster, sessionID string, call *ClientToolCall, result mcp.CallResult) error { + if poster == nil { + return nil + } + return poster.PostAiChatToolResult(ctx, api.ToolResult{ + SessionID: sessionID, + CallID: call.CallID, + Output: mcp.TruncateOutput(result.Output), + IsError: result.IsError, + }) +} + +func localToolConfirmation(ref mcp.ToolRef, call *ClientToolCall) *rootly.AiChatStreamConfirmationRequired { + confirmation := &rootly.AiChatStreamConfirmationRequired{ + Question: "Run local tool " + ref.ServerName + " · " + ref.Tool + "? It can change things on this machine.", + } + if call == nil || len(call.Arguments) == 0 { + return confirmation + } + if ref.ServerID == "builtin" && call.Name == "local_bash" { + if command, ok := call.Arguments["command"].(string); ok { + confirmation.Question = "Run local bash command? It can change things on this machine." + confirmation.Preview = []byte(mcp.MaskApprovalCommand(command)) + return confirmation + } + } + if ref.ServerID == "builtin" && strings.HasPrefix(call.Name, "local_") { + if rawArgs, ok := mcp.MaskApprovalArguments(call.Arguments)["args"]; ok { + if args := previewArgs(rawArgs); args != nil { + command := strings.ReplaceAll(ref.Tool, "_", "-") + confirmation.Question = "Run local CLI " + command + "? It can change things on this machine." + confirmation.Preview = []byte("$ " + command + " " + strings.Join(args, " ")) + return confirmation + } + } + } + if preview, err := json.MarshalIndent(mcp.MaskApprovalArguments(call.Arguments), "", " "); err == nil { + confirmation.Preview = preview + } + return confirmation +} + +func previewArgs(value any) []string { + switch args := value.(type) { + case []string: + return args + case []any: + result := make([]string, len(args)) + for index, arg := range args { + text, ok := arg.(string) + if !ok { + return nil + } + result[index] = text + } + return result + default: + return nil + } +} + +func renderLocalToolCard(ref mcp.ToolRef, call *ClientToolCall, options mcp.ApprovalOptions, width int, noColor bool) string { + return renderConfirmationCardWithKeys(localToolConfirmation(ref, call), width, noColor, time.Now(), localApprovalKeys(options)) +} + +func localApprovalKeys(options mcp.ApprovalOptions) string { + switch { + case len(options.Prefix) > 0 && options.Persist: + prefix := sanitizeTerminalText(strings.Join(options.Prefix, " ")) + return "y once · a allow `" + prefix + "` this session · p always allow `" + prefix + "` · n/Esc cancel" + case options.Session: + return "y once · a allow this tool this session · n/Esc cancel" + } + return confirmationKeys +} diff --git a/internal/agent/localtools_test.go b/internal/agent/localtools_test.go new file mode 100644 index 0000000..56a4bd2 --- /dev/null +++ b/internal/agent/localtools_test.go @@ -0,0 +1,320 @@ +package agent + +import ( + "bytes" + "context" + "encoding/json" + "fmt" + "io" + "net/http" + "net/http/httptest" + "strings" + "sync" + "testing" + "time" + + rootly "github.com/rootlyhq/rootly-go" + + "github.com/rootlyhq/rootly-cli/internal/api" + "github.com/rootlyhq/rootly-cli/internal/config" + "github.com/rootlyhq/rootly-cli/internal/mcp" +) + +const localTestSessionID = "39fd7596-152c-43d3-afce-48b792ae386d" + +type fakeLocalTools struct { + mu sync.Mutex + tools map[string]mcp.ToolRef + calls []string + output string +} + +func (f *fakeLocalTools) ClientTools() (tools []mcp.ClientTool, omitted int) { + out := []mcp.ClientTool{} + for name, ref := range f.tools { + out = append(out, mcp.ClientTool{Name: name, Description: ref.Tool, InputSchema: json.RawMessage(`{"type":"object"}`), ReadOnly: ref.ReadOnly}) + } + return out, 0 +} + +func (f *fakeLocalTools) Lookup(name string) (mcp.ToolRef, bool) { + ref, ok := f.tools[name] + return ref, ok +} + +func (f *fakeLocalTools) NeedsApproval(name string, _ map[string]any) bool { + ref, ok := f.tools[name] + return !ok || !ref.ReadOnly +} + +func (f *fakeLocalTools) Call(_ context.Context, name string, arguments map[string]any) mcp.CallResult { + f.mu.Lock() + defer f.mu.Unlock() + f.calls = append(f.calls, name) + if f.output != "" { + return mcp.CallResult{Output: f.output, IsError: true} + } + return mcp.CallResult{Output: fmt.Sprintf("ran %s with %v", name, arguments["path"])} +} + +func TestLocalToolConfirmationPreservesExecutableCode(t *testing.T) { + ref := mcp.ToolRef{ServerID: "srv", ServerName: "third-party", Tool: "run"} + command := `TOKEN="$(rm -rf ./data)"` + confirmation := localToolConfirmation(ref, &ClientToolCall{ + Name: "third_party_run", + Arguments: map[string]any{ + "command": command, + "password": "x${y}", + }, + }) + var preview map[string]any + if err := json.Unmarshal(confirmation.Preview, &preview); err != nil { + t.Fatalf("decode MCP approval preview: %v", err) + } + if preview["command"] != command || preview["password"] != "x${y}" { + t.Fatalf("MCP approval preview masked executable code: %#v", preview) + } + + confirmation = localToolConfirmation(ref, &ClientToolCall{ + Name: "third_party_run", + Arguments: map[string]any{"password": "hunter22"}, + }) + if err := json.Unmarshal(confirmation.Preview, &preview); err != nil { + t.Fatalf("decode plain-secret preview: %v", err) + } + if preview["password"] != mcp.Mask { + t.Fatalf("MCP approval preview did not mask plain secret: %#v", preview) + } + + confirmation = localToolConfirmation(mcp.ToolRef{ServerID: "builtin", Tool: "kubectl"}, &ClientToolCall{ + Name: "local_kubectl", + Arguments: map[string]any{"args": []any{"--token", "a$(b)"}}, + }) + if !strings.Contains(string(confirmation.Preview), "a$(b)") { + t.Fatalf("CLI approval preview masked executable code: %q", confirmation.Preview) + } + + for _, command := range []string{`export TOKEN="<(rm x)"`, `export TOKEN=">(rm x)"`} { + confirmation = localToolConfirmation(mcp.ToolRef{ServerID: "builtin"}, &ClientToolCall{ + Name: "local_bash", + Arguments: map[string]any{"command": command}, + }) + if !strings.Contains(string(confirmation.Preview), command) { + t.Errorf("bash approval preview masked executable code: %q", confirmation.Preview) + } + } +} + +type toolRoundTrip struct { + server *httptest.Server + streamBody map[string]any + results chan api.ToolResult + accepted int + toolName string + streamCount int + mu sync.Mutex +} + +func newToolRoundTrip(t *testing.T, accepted int, toolName string) *toolRoundTrip { + t.Helper() + trip := &toolRoundTrip{results: make(chan api.ToolResult, 1), accepted: accepted, toolName: toolName} + trip.server = httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + switch r.URL.Path { + case "/v1/ai/chat/stream": + body, _ := io.ReadAll(r.Body) + trip.mu.Lock() + trip.streamCount++ + _ = json.Unmarshal(body, &trip.streamBody) + trip.mu.Unlock() + w.Header().Set("Content-Type", "text/event-stream") + flusher := w.(http.Flusher) + _, _ = fmt.Fprintf(w, "event: session_id\ndata: {\"session_id\":%q,\"client_tools_accepted\":%d}\n\n", localTestSessionID, trip.accepted) + if trip.toolName != "" { + _, _ = fmt.Fprintf(w, "event: client_tool_call\ndata: {\"call_id\":\"call-1\",\"name\":%q,\"arguments\":{\"path\":\"/tmp/a\"}}\n\n", trip.toolName) + flusher.Flush() + select { + case result := <-trip.results: + trip.results <- result + _, _ = fmt.Fprintf(w, "event: text\ndata: {\"content\":\"got %s\"}\n\n", strings.ReplaceAll(result.Output, `"`, ``)) + case <-time.After(5 * time.Second): + _, _ = fmt.Fprint(w, "event: text\ndata: {\"content\":\"timed out\"}\n\n") + } + } else { + _, _ = fmt.Fprint(w, "event: text\ndata: {\"content\":\"hello\"}\n\n") + } + _, _ = fmt.Fprint(w, "event: done\ndata: {\"status\":\"complete\"}\n\n") + case "/v1/ai/chat/tool_results": + if r.Method != http.MethodPost || r.Header.Get("Authorization") != "Bearer test-token" { + t.Errorf("tool_results request = %s auth=%q", r.Method, r.Header.Get("Authorization")) + } + var result api.ToolResult + if err := json.NewDecoder(r.Body).Decode(&result); err != nil { + t.Errorf("decode tool result: %v", err) + } + trip.results <- result + w.WriteHeader(http.StatusAccepted) + default: + t.Errorf("unexpected request %s", r.URL.Path) + w.WriteHeader(http.StatusNotFound) + } + })) + t.Cleanup(trip.server.Close) + return trip +} + +func (trip *toolRoundTrip) run(t *testing.T, tools LocalTools, stdinTTY bool, stdin string) (stdout, stderr string, result api.ToolResult) { + t.Helper() + client, err := api.NewClient(&config.Config{APIKey: "test-token", Endpoint: trip.server.URL}) + if err != nil { + t.Fatal(err) + } + var outBuf, errBuf bytes.Buffer + options := newPrintOptions(&outBuf, &errBuf) + options.Stdin = strings.NewReader(stdin) + options.IsStdinTTY = func() bool { return stdinTTY } + options.IsStderrTTY = func() bool { return stdinTTY } + options.NoColor = true + options.LocalTools = tools + options.ToolResults = client + ctx, cancel := context.WithTimeout(context.Background(), 10*time.Second) + defer cancel() + if err := RunPrint(ctx, client, options); err != nil { + t.Fatalf("RunPrint() error = %v (stderr %q)", err, errBuf.String()) + } + select { + case result = <-trip.results: + default: + } + return outBuf.String(), errBuf.String(), result +} + +func TestPrintClientToolCallRoundTripReadOnly(t *testing.T) { + tools := &fakeLocalTools{tools: map[string]mcp.ToolRef{"local_fs_read_file": {ServerID: "user:fs", ServerName: "fs", Tool: "read_file", ReadOnly: true}}} + trip := newToolRoundTrip(t, 1, "local_fs_read_file") + stdout, stderr, result := trip.run(t, tools, true, "") + + sent, _ := trip.streamBody["client_tools"].([]any) + if len(sent) != 1 || sent[0].(map[string]any)["name"] != "local_fs_read_file" || sent[0].(map[string]any)["read_only"] != true { + t.Fatalf("client_tools body = %v", trip.streamBody) + } + if result.SessionID != localTestSessionID || result.CallID != "call-1" || result.IsError || result.Output != "ran local_fs_read_file with /tmp/a" { + t.Fatalf("result = %+v", result) + } + if !strings.Contains(stdout, "got ran local_fs_read_file") { + t.Fatalf("stdout = %q", stdout) + } + if !strings.Contains(stderr, "⚙ local · fs · read_file") || strings.Contains(stderr, localToolsOffMessage) { + t.Fatalf("stderr = %q", stderr) + } +} + +func TestPrintLocalToolResultControlsCompletionActivity(t *testing.T) { + for _, test := range []struct { + name string + output string + wantCross int + wantCheck int + }{ + {name: "error", output: "path is denied", wantCross: 1}, + {name: "success", wantCheck: 1}, + } { + t.Run(test.name, func(t *testing.T) { + tools := &fakeLocalTools{ + tools: map[string]mcp.ToolRef{"local_fs_read_file": {ServerID: "user:fs", ServerName: "fs", Tool: "read_file", ReadOnly: true}}, + output: test.output, + } + var stdout, stderr bytes.Buffer + options := newPrintOptions(&stdout, &stderr) + options.IsStderrTTY = func() bool { return true } + options.OutputFormat = "text" + options.LocalTools = tools + run := &printRun{options: options, ctx: context.Background()} + run.activities.Handle(activityStatusEvent("Local read file...")) + run.handleClientToolCall(&ClientToolCall{CallID: "call-1", Name: "local_fs_read_file"}) + if test.wantCross > 0 { + run.writeActivityLines(run.activities.Handle(Event{ + Type: rootly.AiChatStreamEventTaskUpdate, + TaskUpdate: &rootly.AiChatStreamTaskUpdate{ + ID: "local-read", + Title: "Local read file failed: path is denied", + Status: "failed", + }, + })) + } else { + run.writeActivityLines(run.activities.Finish(true)) + } + if got := strings.Count(stderr.String(), "✗"); got != test.wantCross { + t.Errorf("failure markers = %d, want %d:\n%s", got, test.wantCross, stderr.String()) + } + if got := strings.Count(stderr.String(), "✓"); got != test.wantCheck { + t.Errorf("success markers = %d, want %d:\n%s", got, test.wantCheck, stderr.String()) + } + }) + } +} + +func TestPrintClientToolWriteDeclinedWithoutTTY(t *testing.T) { + tools := &fakeLocalTools{tools: map[string]mcp.ToolRef{"local_fs_write_file": {ServerID: "user:fs", ServerName: "fs", Tool: "write_file"}}} + trip := newToolRoundTrip(t, 1, "local_fs_write_file") + _, _, result := trip.run(t, tools, false, "y\n") + if !result.IsError || result.Output != "User declined" { + t.Fatalf("result = %+v", result) + } + if len(tools.calls) != 0 { + t.Fatalf("declined tool must not run, calls = %v", tools.calls) + } +} + +func TestPrintClientToolWriteApprovedOnTTY(t *testing.T) { + tools := &fakeLocalTools{tools: map[string]mcp.ToolRef{"local_fs_write_file": {ServerID: "user:fs", ServerName: "fs", Tool: "write_file"}}} + trip := newToolRoundTrip(t, 1, "local_fs_write_file") + _, _, result := trip.run(t, tools, true, "y\n") + if result.IsError || len(tools.calls) != 1 { + t.Fatalf("result = %+v calls = %v", result, tools.calls) + } + + declined := &fakeLocalTools{tools: tools.tools} + trip = newToolRoundTrip(t, 1, "local_fs_write_file") + _, _, result = trip.run(t, declined, true, "n\n") + if !result.IsError || result.Output != "User declined" || len(declined.calls) != 0 { + t.Fatalf("result = %+v calls = %v", result, declined.calls) + } +} + +func TestPrintUnknownClientToolReturnsError(t *testing.T) { + tools := &fakeLocalTools{tools: map[string]mcp.ToolRef{"local_fs_read_file": {ServerName: "fs", Tool: "read_file", ReadOnly: true}}} + trip := newToolRoundTrip(t, 1, "local_other_tool") + _, _, result := trip.run(t, tools, false, "") + if !result.IsError || result.CallID != "call-1" { + t.Fatalf("result = %+v", result) + } +} + +func TestPrintLocalToolsOffForOrg(t *testing.T) { + tools := &fakeLocalTools{tools: map[string]mcp.ToolRef{"local_fs_read_file": {ServerName: "fs", Tool: "read_file", ReadOnly: true}}} + trip := newToolRoundTrip(t, 0, "") + stdout, stderr, _ := trip.run(t, tools, false, "") + if strings.Count(stderr, localToolsOffMessage) != 1 || stdout != "hello\n" { + t.Fatalf("stdout = %q stderr = %q", stdout, stderr) + } + + none := &fakeLocalTools{tools: map[string]mcp.ToolRef{}} + trip = newToolRoundTrip(t, 0, "") + _, stderr, _ = trip.run(t, none, false, "") + if strings.Contains(stderr, localToolsOffMessage) { + t.Fatalf("no tools sent, notice must not show: %q", stderr) + } + if _, ok := trip.streamBody["client_tools"]; ok { + t.Fatalf("no client_tools key expected: %v", trip.streamBody) + } +} + +func TestDecodeClientToolCallRequiresIDs(t *testing.T) { + if _, _, err := decodeEvent(eventOf(EventClientToolCall, `{"name":"local_x"}`)); err == nil { + t.Fatal("missing call_id should fail") + } + event, _, err := decodeEvent(eventOf(EventClientToolCall, `{"call_id":"c","name":"local_x","arguments":{"a":1}}`)) + if err != nil || event.ClientToolCall == nil || event.ClientToolCall.Arguments["a"] != float64(1) { + t.Fatalf("event = %+v err = %v", event, err) + } +} diff --git a/internal/agent/markdown_links.go b/internal/agent/markdown_links.go new file mode 100644 index 0000000..748075f --- /dev/null +++ b/internal/agent/markdown_links.go @@ -0,0 +1,451 @@ +package agent + +import ( + "bytes" + "net/url" + "sort" + "strings" + + "charm.land/lipgloss/v2" + "github.com/charmbracelet/x/ansi" + "github.com/yuin/goldmark" + "github.com/yuin/goldmark/ast" + goldmarktext "github.com/yuin/goldmark/text" +) + +type markdownLink struct { + destination string + host string + hyperlink bool +} + +type markdownLinkReplacement struct { + start int + end int + text string + link markdownLink +} + +func terminalMarkdown(markdown string, noColor bool) string { + markdown = sanitizeTerminalText(markdown) + if noColor { + return markdown + } + return preprocessMarkdownLinks(markdown) +} + +func preprocessMarkdownLinks(markdown string) string { + processed, _ := preprocessMarkdownLinksWithURLs(markdown, "") + return processed +} + +func preprocessMarkdownLinksWithURLs(markdown, apiHost string) (processed string, links []markdownLink) { + markdown = sanitizeTerminalText(markdown) + source := []byte(markdown) + document := goldmark.New().Parser().Parse(goldmarktext.NewReader(source)) + replacements := make([]markdownLinkReplacement, 0) + + _ = ast.Walk(document, func(node ast.Node, entering bool) (ast.WalkStatus, error) { + if !entering { + return ast.WalkContinue, nil + } + link, ok := node.(*ast.Link) + if !ok || len(link.Destination) == 0 { + return ast.WalkContinue, nil + } + label := markdownLinkLabel(link, source) + destination := string(link.Destination) + if label == "" || label == destination { + return ast.WalkContinue, nil + } + labelStart, labelEnd, ok := markdownLinkTextRange(link, source) + if !ok { + return ast.WalkContinue, nil + } + start, closingBracket, end, ok := findInlineLinkRange(source, labelStart, labelEnd, destination) + if !ok { + return ast.WalkContinue, nil + } + replacement := "\uE000" + string(source[start+1:closingBracket]) + "\uE001" + safeLink := newMarkdownLink(destination, apiHost) + replacements = append(replacements, markdownLinkReplacement{ + start: start, + end: end, + text: replacement, + link: safeLink, + }) + return ast.WalkContinue, nil + }) + + if len(replacements) == 0 { + return markdown, nil + } + sort.Slice(replacements, func(i, j int) bool { + return replacements[i].start < replacements[j].start + }) + links = make([]markdownLink, len(replacements)) + for index, replacement := range replacements { + links[index] = replacement.link + } + for index := len(replacements) - 1; index >= 0; index-- { + replacement := replacements[index] + source = append(source[:replacement.start], append([]byte(replacement.text), source[replacement.end:]...)...) + } + return string(source), links +} + +func rewriteDashboardLinks(markdown, webBase, sessionID string) string { + if !strings.Contains(strings.ToLower(markdown), "#ai-add-to-dashboard") { + return markdown + } + destination := strings.TrimRight(webBase, "/") + "/account" + if sessionID != "" { + destination += "?chat=" + url.QueryEscape(sessionID) + } + + source := []byte(markdown) + document := goldmark.New().Parser().Parse(goldmarktext.NewReader(source)) + replacements := make([]markdownLinkReplacement, 0) + _ = ast.Walk(document, func(node ast.Node, entering bool) (ast.WalkStatus, error) { + if !entering { + return ast.WalkContinue, nil + } + link, ok := node.(*ast.Link) + if !ok { + return ast.WalkContinue, nil + } + href := strings.ToLower(string(link.Destination)) + if !strings.HasPrefix(href, "#ai-add-to-dashboard:") && + !strings.HasPrefix(href, "#ai-add-to-dashboard%3a") { + return ast.WalkContinue, nil + } + labelStart, labelEnd, ok := markdownLinkTextRange(link, source) + if !ok { + return ast.WalkContinue, nil + } + start, _, end, ok := findInlineLinkRange(source, labelStart, labelEnd, string(link.Destination)) + if !ok { + return ast.WalkContinue, nil + } + replacements = append(replacements, markdownLinkReplacement{ + start: start, + end: end, + text: "[Add to dashboard: open in web](" + destination + ")", + }) + return ast.WalkContinue, nil + }) + for index := len(replacements) - 1; index >= 0; index-- { + replacement := replacements[index] + source = append(source[:replacement.start], append([]byte(replacement.text), source[replacement.end:]...)...) + } + return string(source) +} + +func newMarkdownLink(destination, apiHost string) markdownLink { + link := markdownLink{destination: sanitizeTerminalText(destination)} + parsed, err := url.Parse(link.destination) + if err != nil { + return link + } + scheme := strings.ToLower(parsed.Scheme) + if scheme != "http" && scheme != "https" { + return link + } + host := strings.TrimSuffix(strings.ToLower(parsed.Hostname()), ".") + if host == "" { + return link + } + link.hyperlink = true + domain := rootlyAppDomain(apiHost) + if domain != "" && (host == domain || strings.HasSuffix(host, "."+domain)) { + return link + } + link.host = host + return link +} + +func rootlyAppDomain(apiHost string) string { + apiHost = strings.TrimSpace(sanitizeTerminalText(apiHost)) + if apiHost == "" { + return "" + } + if !strings.Contains(apiHost, "://") { + apiHost = "//" + apiHost + } + parsed, err := url.Parse(apiHost) + if err != nil { + return "" + } + domain := strings.TrimSuffix(strings.ToLower(parsed.Hostname()), ".") + return strings.TrimPrefix(domain, "api.") +} + +func restoreMarkdownLinks(rendered string, links []markdownLink) string { + const ( + startMarker = "\uE000" + endMarker = "\uE001" + ) + var output strings.Builder + var segment strings.Builder + linkIndex := 0 + currentLink := -1 + for rendered != "" { + startIndex := strings.Index(rendered, startMarker) + endIndex := strings.Index(rendered, endMarker) + if startIndex < 0 && endIndex < 0 { + break + } + markerIndex := startIndex + marker := startMarker + if markerIndex < 0 || (endIndex >= 0 && endIndex < markerIndex) { + markerIndex = endIndex + marker = endMarker + } + if currentLink >= 0 { + segment.WriteString(rendered[:markerIndex]) + } else { + output.WriteString(rendered[:markerIndex]) + } + if marker == startMarker { + if currentLink >= 0 { + output.WriteString(segment.String()) + segment.Reset() + } + currentLink = linkIndex + linkIndex++ + } else { + if currentLink >= 0 && currentLink < len(links) { + link := links[currentLink] + if link.hyperlink { + output.WriteString(ansi.SetHyperlink(link.destination)) + output.WriteString(styleMarkdownLinkText(segment.String())) + output.WriteString(ansi.ResetHyperlink()) + if link.host != "" { + output.WriteString(lipgloss.NewStyle().Faint(true).Render(" (" + link.host + ")")) + } + } else { + output.WriteString(segment.String()) + output.WriteString(" <" + link.destination + ">") + } + } else { + output.WriteString(segment.String()) + } + currentLink = -1 + segment.Reset() + } + rendered = rendered[markerIndex+len(marker):] + } + if currentLink >= 0 { + output.WriteString(segment.String()) + } + output.WriteString(rendered) + return output.String() +} + +func styleMarkdownLinkText(segment string) string { + style := lipgloss.NewStyle(). + Underline(true). + Foreground(lipgloss.Color(metricChartColor)) + lines := strings.Split(ansi.Strip(segment), "\n") + for index, line := range lines { + text := strings.TrimLeft(line, " ") + indent := line[:len(line)-len(text)] + if text == "" { + lines[index] = indent + continue + } + lines[index] = indent + style.Render(text) + } + return strings.Join(lines, "\n") +} + +func markdownLinkLabel(link ast.Node, source []byte) string { + var label strings.Builder + _ = ast.Walk(link, func(node ast.Node, entering bool) (ast.WalkStatus, error) { + if !entering { + return ast.WalkContinue, nil + } + if text, ok := node.(*ast.Text); ok { + label.Write(text.Value(source)) + } + return ast.WalkContinue, nil + }) + return label.String() +} + +func markdownLinkTextRange(link ast.Node, source []byte) (start, end int, found bool) { + start, end = len(source), 0 + _ = ast.Walk(link, func(node ast.Node, entering bool) (ast.WalkStatus, error) { + if !entering { + return ast.WalkContinue, nil + } + text, ok := node.(*ast.Text) + if !ok || text.Segment.Len() == 0 { + return ast.WalkContinue, nil + } + start = min(start, text.Segment.Start) + end = max(end, text.Segment.Stop) + found = true + return ast.WalkContinue, nil + }) + return start, end, found +} + +func findInlineLinkRange(source []byte, labelStart, labelEnd int, destination string) (start, closingBracket, end int, found bool) { + for start := labelStart - 1; start >= 0; start-- { + if source[start] != '[' || escapedMarkdownCharacter(source, start) { + continue + } + closingBracket := matchingMarkdownBracket(source, start) + if closingBracket < labelEnd || closingBracket+1 >= len(source) || source[closingBracket+1] != '(' { + continue + } + closingParen, ok := matchingMarkdownParen(source, closingBracket+1) + if !ok { + continue + } + rawDestination := markdownLinkDestination(source[closingBracket+2 : closingParen]) + if rawDestination == destination { + return start, closingBracket, closingParen + 1, true + } + } + return 0, 0, 0, false +} + +func matchingMarkdownBracket(source []byte, start int) int { + depth := 1 + for index := start + 1; index < len(source); index++ { + if escapedMarkdownCharacter(source, index) { + continue + } + switch source[index] { + case '[': + depth++ + case ']': + depth-- + if depth == 0 { + return index + } + } + } + return -1 +} + +func matchingMarkdownParen(source []byte, start int) (int, bool) { + depth := 1 + var quote byte + angle := false + for index := start + 1; index < len(source); index++ { + character := source[index] + if escapedMarkdownCharacter(source, index) { + continue + } + if quote != 0 { + if character == quote { + quote = 0 + } + continue + } + if character == '"' || character == '\'' { + quote = character + continue + } + if character == '<' { + angle = true + continue + } + if character == '>' { + angle = false + continue + } + if angle { + continue + } + switch character { + case '(': + depth++ + case ')': + depth-- + if depth == 0 { + return index, true + } + } + } + return 0, false +} + +func trimIncompleteTrailingMarkdownLink(markdown string) string { + source := []byte(markdown) + incompleteStart := -1 + for index := 0; index < len(source); index++ { + if source[index] != '[' || escapedMarkdownCharacter(source, index) { + continue + } + closingBracket := matchingMarkdownBracket(source, index) + if closingBracket < 0 { + incompleteStart = index + continue + } + if closingBracket+1 >= len(source) || source[closingBracket+1] != '(' { + continue + } + closingParen, complete := matchingMarkdownParen(source, closingBracket+1) + if !complete { + incompleteStart = index + continue + } + index = closingParen + } + if incompleteStart >= 0 { + return string(source[:incompleteStart]) + } + return markdown +} + +func markdownLinkDestination(raw []byte) string { + raw = bytes.TrimSpace(raw) + if len(raw) > 0 && raw[0] == '<' { + if end := bytes.IndexByte(raw, '>'); end >= 0 { + raw = raw[1:end] + } + } else { + depth := 0 + for index := 0; index < len(raw); index++ { + if escapedMarkdownCharacter(raw, index) { + continue + } + switch raw[index] { + case '(': + depth++ + case ')': + if depth == 0 { + raw = raw[:index] + index = len(raw) + } else { + depth-- + } + case ' ', '\t', '\r', '\n': + if depth == 0 { + raw = raw[:index] + index = len(raw) + } + } + } + } + var destination strings.Builder + for index := 0; index < len(raw); index++ { + if raw[index] == '\\' && index+1 < len(raw) { + index++ + } + destination.WriteByte(raw[index]) + } + return destination.String() +} + +func escapedMarkdownCharacter(source []byte, index int) bool { + backslashes := 0 + for index--; index >= 0 && source[index] == '\\'; index-- { + backslashes++ + } + return backslashes%2 == 1 +} diff --git a/internal/agent/markdown_tables.go b/internal/agent/markdown_tables.go new file mode 100644 index 0000000..d55941d --- /dev/null +++ b/internal/agent/markdown_tables.go @@ -0,0 +1,149 @@ +package agent + +import ( + "strings" + + "github.com/charmbracelet/x/ansi" +) + +func bulletizeWideTables(markdown string, width int) string { + if width <= 0 || !strings.Contains(markdown, "|") { + return markdown + } + lines := strings.Split(markdown, "\n") + output := make([]string, 0, len(lines)) + fence := "" + for index := 0; index < len(lines); index++ { + line := lines[index] + if fence != "" { + output = append(output, line) + if closesFence(line, fence) { + fence = "" + } + continue + } + if marker, ok := openingMarkdownFence(line); ok { + fence = marker + output = append(output, line) + continue + } + if index+1 < len(lines) && isTableRow(line) && isTableSeparator(lines[index+1]) { + end := index + 2 + for end < len(lines) && isTableRow(lines[end]) { + end++ + } + table := lines[index:end] + if len(table) > 2 && markdownTableWidth(table) > width { + output = append(output, tableAsBullets(table)...) + } else { + output = append(output, table...) + } + index = end - 1 + continue + } + output = append(output, line) + } + return strings.Join(output, "\n") +} + +func closesFence(line, fence string) bool { + trimmed := strings.TrimLeft(line, " ") + if len(line)-len(trimmed) > 3 { + return false + } + trimmed = strings.TrimSpace(trimmed) + return strings.HasPrefix(trimmed, fence) && strings.Trim(trimmed, fence[:1]) == "" +} + +func isTableRow(line string) bool { + trimmed := strings.TrimSpace(line) + return strings.Count(trimmed, "|") >= 1 && strings.Contains(trimmed, "|") && trimmed != "" +} + +func isTableSeparator(line string) bool { + cells := tableCells(line) + if len(cells) == 0 { + return false + } + for _, cell := range cells { + cell = strings.Trim(cell, ":") + if cell == "" || strings.Trim(cell, "-") != "" { + return false + } + } + return true +} + +func tableCells(line string) []string { + trimmed := strings.TrimSpace(line) + trimmed = strings.TrimPrefix(trimmed, "|") + if withoutPipe, ok := strings.CutSuffix(trimmed, "|"); ok { + if backslashes := len(withoutPipe) - len(strings.TrimRight(withoutPipe, `\`)); backslashes%2 == 0 { + trimmed = withoutPipe + } + } + var cells []string + var cell strings.Builder + for index := 0; index < len(trimmed); index++ { + switch { + case trimmed[index] == '\\' && index+1 < len(trimmed): + if trimmed[index+1] != '|' { + cell.WriteByte('\\') + } + cell.WriteByte(trimmed[index+1]) + index++ + case trimmed[index] == '|': + cells = append(cells, strings.TrimSpace(cell.String())) + cell.Reset() + default: + cell.WriteByte(trimmed[index]) + } + } + return append(cells, strings.TrimSpace(cell.String())) +} + +func markdownTableWidth(table []string) int { + widths := []int{} + for index, line := range table { + if index == 1 { + continue + } + for column, cell := range tableCells(line) { + if column >= len(widths) { + widths = append(widths, 0) + } + widths[column] = max(widths[column], ansi.StringWidth(cell)) + } + } + total := 1 + for _, width := range widths { + total += width + 3 + } + return total +} + +func tableAsBullets(table []string) []string { + headers := tableCells(table[0]) + bullets := make([]string, 0, len(table)-2) + for _, line := range table[2:] { + cells := tableCells(line) + parts := make([]string, 0, len(cells)) + for column, cell := range cells { + if cell == "" { + continue + } + if column == 0 { + parts = append(parts, "**"+cell+"**") + continue + } + if column < len(headers) && headers[column] != "" { + cell = headers[column] + ": " + cell + } + parts = append(parts, cell) + } + if len(parts) > 0 { + bullets = append(bullets, "- "+strings.Join(parts, " · ")) + } + } + return bullets +} diff --git a/internal/agent/market.go b/internal/agent/market.go new file mode 100644 index 0000000..a58b77b --- /dev/null +++ b/internal/agent/market.go @@ -0,0 +1,209 @@ +package agent + +import ( + "context" + "fmt" + "strings" + + tea "charm.land/bubbletea/v2" + "charm.land/lipgloss/v2" + "github.com/charmbracelet/x/ansi" + + "github.com/rootlyhq/rootly-cli/internal/mcp" +) + +const marketResultLimit = 50 + +var loadMarketCatalog = func(ctx context.Context, refresh bool) ([]mcp.MarketEntry, error) { + return mcp.LoadCatalog(ctx, nil, mcp.RegistryURL, mcp.IndexPath(), refresh, nil) +} + +type marketState struct { + query string + loading bool + err string + catalog []mcp.MarketEntry + results []mcp.MarketEntry + cursor int + offset int + preview *mcp.MarketEntry + request int +} + +type marketLoadedMsg struct { + request int + catalog []mcp.MarketEntry + err error +} + +func (m *TUIModel) openMarket() tea.Cmd { + if m.mcp == nil { + m.mcpPage.message = "Local MCP is unavailable in this session" + return nil + } + m.mcpPage.market = &marketState{} + return m.loadMarket(false) +} + +func (m *TUIModel) loadMarket(refresh bool) tea.Cmd { + market := m.mcpPage.market + m.marketRequest++ + market.request = m.marketRequest + market.loading, market.err = true, "" + request, ctx := market.request, m.ctx + return func() tea.Msg { + catalog, err := loadMarketCatalog(ctx, refresh) + return marketLoadedMsg{request: request, catalog: catalog, err: err} + } +} + +func (m *TUIModel) handleMarketLoaded(msg marketLoadedMsg) { + market := m.mcpPage.market + if market == nil || msg.request != market.request { + return + } + market.loading, market.catalog = false, msg.catalog + if msg.err != nil { + market.err = sanitizeTerminalText(msg.err.Error()) + } + m.searchMarket() +} + +func (m *TUIModel) searchMarket() { + market := m.mcpPage.market + market.results = mcp.SearchMarket(market.catalog, market.query, marketResultLimit) + market.cursor, market.offset = 0, 0 +} + +func (m *TUIModel) handleMarketKey(msg tea.KeyPressMsg) tea.Cmd { + market := m.mcpPage.market + key := msg.String() + if market.preview != nil { + switch key { + case "y": + m.addFromMarket(*market.preview) + case "n", "esc", "backspace": + market.preview = nil + } + return nil + } + switch key { + case "esc", "ctrl+c": + m.mcpPage.market = nil + case "up": + market.cursor = max(0, market.cursor-1) + case "down": + market.cursor = min(max(0, len(market.results)-1), market.cursor+1) + case "enter": + if market.cursor < len(market.results) { + entry := market.results[market.cursor] + market.preview = &entry + } + case "ctrl+r": + return m.loadMarket(true) + case "backspace": + if market.query != "" { + runes := []rune(market.query) + market.query = string(runes[:len(runes)-1]) + m.searchMarket() + } + default: + if msg.Text != "" { + market.query += msg.Text + m.searchMarket() + } + } + return nil +} + +func (m *TUIModel) addFromMarket(entry mcp.MarketEntry) { + server, err := m.mcp.AddFromMarket(entry) + if err != nil { + m.mcpPage.market.preview = nil + m.mcpPage.market.err = sanitizeTerminalText(err.Error()) + return + } + m.mcpPage.market = nil + m.mcpPage.message = fmt.Sprintf("Added %s to ~/.rootly-cli/mcp.json (off). Select it and press space to turn it on", sanitizeTerminalText(server.Name)) + m.refreshMCPViews() + for index, row := range m.mcpRows() { + if row.kind == mcpRowServer && row.view.Server.ID() == server.ID() { + m.mcpPage.cursor = index + } + } +} + +func (m *TUIModel) marketView() string { + market := m.mcpPage.market + width := max(20, m.width-1) + bold := lipgloss.NewStyle().Bold(true) + if market.preview != nil { + return m.marketPreviewView(*market.preview, width, bold) + } + lines := []string{m.mcpStyle(bold, "MCP marketplace"), faint(m.noColor, "Official MCP registry + curated list · Rootly-integrated vendors first"), "Search: " + sanitizeTerminalText(market.query) + "▌"} + switch { + case market.loading && len(market.catalog) == 0: + lines = append(lines, "", faint(m.noColor, "Indexing the MCP registry (cached in ~/.rootly-cli/mcp-index.json)…")) + case market.err != "": + lines = append(lines, faint(m.noColor, "✘ "+market.err)) + default: + lines = append(lines, "") + } + height := max(1, m.mcpListHeight()-2) + if market.cursor < market.offset { + market.offset = market.cursor + } + if market.cursor >= market.offset+height { + market.offset = market.cursor - height + 1 + } + end := min(len(market.results), market.offset+height) + for index := market.offset; index < end; index++ { + lines = append(lines, ansi.Truncate(m.renderMarketRow(market.results[index], index == market.cursor), width, "…")) + } + if len(market.results) == 0 && !market.loading { + lines = append(lines, faint(m.noColor, "No MCP servers match")) + } + for len(lines) < height+4 { + lines = append(lines, "") + } + lines = append(lines, "", faint(m.noColor, ansi.Truncate("type to search · ↑/↓ move · Enter preview · ctrl+r refresh index · Esc back", width, "…"))) + return strings.Join(lines, "\n") +} + +func (m *TUIModel) renderMarketRow(entry mcp.MarketEntry, selected bool) string { + marker := " " + if selected { + marker = "› " + } + name := sanitizeTerminalText(entry.Name) + if entry.Rootly { + name += " ★" + } + return marker + fmt.Sprintf("%-24s ", name) + faint(m.noColor, sanitizeTerminalText(strings.Join(strings.Fields(entry.Description), " "))) +} + +func (m *TUIModel) marketPreviewView(entry mcp.MarketEntry, width int, bold lipgloss.Style) string { + lines := []string{m.mcpStyle(bold, "Add "+sanitizeTerminalText(entry.Name)+" to your user MCPs?"), ""} + if entry.Description != "" { + lines = append(lines, sanitizeTerminalText(strings.Join(strings.Fields(entry.Description), " ")), "") + } + verb := "Runs" + if entry.Config.URL != "" { + verb = "Connects to" + } + lines = append(lines, + verb+": "+sanitizeTerminalText(mcp.InstallLabel(entry)), + "Writes to: ~/.rootly-cli/mcp.json (user scope)", + ) + if entry.Author != "" { + lines = append(lines, "Publisher: "+sanitizeTerminalText(entry.Author)) + } + if entry.EnvHint != "" { + lines = append(lines, "Needs: "+sanitizeTerminalText(entry.EnvHint)) + } + lines = append(lines, "", "It is added switched off: nothing runs until you turn it on in /mcp.", "", faint(m.noColor, "y add · n back")) + for index, line := range lines { + lines[index] = ansi.Truncate(line, width, "…") + } + return strings.Join(lines, "\n") +} diff --git a/internal/agent/market_test.go b/internal/agent/market_test.go new file mode 100644 index 0000000..d4608b4 --- /dev/null +++ b/internal/agent/market_test.go @@ -0,0 +1,144 @@ +package agent + +import ( + "context" + "strings" + "testing" + + tea "charm.land/bubbletea/v2" + + "github.com/rootlyhq/rootly-cli/internal/mcp" +) + +func stubMarket(t *testing.T) *int { + t.Helper() + calls := 0 + previous := loadMarketCatalog + loadMarketCatalog = func(context.Context, bool) ([]mcp.MarketEntry, error) { + calls++ + return mcp.Catalog([]mcp.MarketEntry{{Name: "pager-mcp", Description: "Paging tool", Author: "acme", Config: mcp.ServerConfig{Command: "npx", Args: []string{"-y", "@acme/pager-mcp"}}}}), nil + } + t.Cleanup(func() { loadMarketCatalog = previous }) + return &calls +} + +func typeKeys(model *TUIModel, text string) { + for _, char := range text { + model.Update(tea.KeyPressMsg{Code: char, Text: string(char)}) + } +} + +func TestMarketplaceSearchPreviewAndAddToUserScopeOff(t *testing.T) { + calls := stubMarket(t) + manager := &fakeMCPManager{views: testViews(), changes: make(chan struct{})} + model := newMCPTestModel(t, manager, &fakeConnectorAPI{}) + model.mcpPage.open = true + _, cmd := model.Update(tea.KeyPressMsg{Code: 'm', Text: "m"}) + runCmd(model, cmd) + if *calls != 1 || model.mcpPage.market == nil { + t.Fatalf("marketplace did not open: calls=%d", *calls) + } + view := model.mcpPageView() + if !strings.Contains(view, "MCP marketplace") || strings.Index(view, "datadog ★") > strings.Index(view, "github") { + t.Fatalf("Rootly vendors should be listed first:\n%s", view) + } + typeKeys(model, "pager") + if got := model.mcpPage.market.results; len(got) == 0 || got[0].Name != "pager-mcp" { + t.Fatalf("results = %+v", got) + } + model.Update(tea.KeyPressMsg{Code: tea.KeyEnter}) + preview := model.mcpPageView() + for _, want := range []string{"Add pager-mcp to your user MCPs?", "Runs: npx -y @acme/pager-mcp", "Writes to: ~/.rootly-cli/mcp.json (user scope)", "nothing runs until you turn it on"} { + if !strings.Contains(preview, want) { + t.Fatalf("preview missing %q:\n%s", want, preview) + } + } + if len(manager.added) != 0 { + t.Fatal("nothing may be added before y") + } + model.Update(tea.KeyPressMsg{Code: 'y', Text: "y"}) + if len(manager.added) != 1 || manager.added[0].Name != "pager-mcp" || model.mcpPage.market != nil { + t.Fatalf("added = %+v market=%v", manager.added, model.mcpPage.market) + } + if !strings.Contains(model.mcpPage.message, "Added pager-mcp") || !strings.Contains(model.mcpPage.message, "(off)") { + t.Fatalf("message = %q", model.mcpPage.message) + } + if row, ok := model.selectedMCPRow(); !ok || row.view.Server.Name != "pager-mcp" || row.view.Status != mcp.StatusOff { + t.Fatalf("cursor should land on the new server, switched off: %+v", row) + } +} + +func TestMarketplaceDeclineAndDuplicate(t *testing.T) { + stubMarket(t) + manager := &fakeMCPManager{views: testViews(), changes: make(chan struct{})} + model := newMCPTestModel(t, manager, &fakeConnectorAPI{}) + model.mcpPage.open = true + _, cmd := model.Update(tea.KeyPressMsg{Code: 'm', Text: "m"}) + runCmd(model, cmd) + model.Update(tea.KeyPressMsg{Code: tea.KeyEnter}) + model.Update(tea.KeyPressMsg{Code: 'n', Text: "n"}) + if model.mcpPage.market.preview != nil || len(manager.added) != 0 { + t.Fatal("n should go back without adding") + } + typeKeys(model, "sentry") + model.Update(tea.KeyPressMsg{Code: tea.KeyEnter}) + model.Update(tea.KeyPressMsg{Code: 'y', Text: "y"}) + if len(manager.added) != 0 || !strings.Contains(model.mcpPage.market.err, "already exists") { + t.Fatalf("duplicate should be refused: added=%v err=%q", manager.added, model.mcpPage.market.err) + } + model.Update(tea.KeyPressMsg{Code: tea.KeyEscape}) + if model.mcpPage.market != nil || !model.mcpPage.open { + t.Fatal("Esc should return to the /mcp list") + } +} + +func TestMarketplaceIgnoresLoadsFromEarlierVisits(t *testing.T) { + stubMarket(t) + manager := &fakeMCPManager{views: testViews(), changes: make(chan struct{})} + model := newMCPTestModel(t, manager, &fakeConnectorAPI{}) + model.mcpPage.open = true + model.openMarket() + first := model.mcpPage.market.request + model.mcpPage.market = nil + model.openMarket() + if model.mcpPage.market.request == first { + t.Fatal("a new marketplace visit must not reuse the previous request id") + } + model.handleMarketLoaded(marketLoadedMsg{request: first, catalog: []mcp.MarketEntry{{Name: "stale"}}}) + if model.mcpPage.market.catalog != nil || !model.mcpPage.market.loading { + t.Fatalf("stale load populated the new visit: %+v", model.mcpPage.market) + } +} + +func TestMarketplacePreviewIgnoresEnter(t *testing.T) { + stubMarket(t) + manager := &fakeMCPManager{views: testViews(), changes: make(chan struct{})} + model := newMCPTestModel(t, manager, &fakeConnectorAPI{}) + model.mcpPage.open = true + _, cmd := model.Update(tea.KeyPressMsg{Code: 'm', Text: "m"}) + runCmd(model, cmd) + typeKeys(model, "pager") + model.Update(tea.KeyPressMsg{Code: tea.KeyEnter}) + model.Update(tea.KeyPressMsg{Code: tea.KeyEnter}) + if len(manager.added) != 0 || model.mcpPage.market == nil || model.mcpPage.market.preview == nil { + t.Fatalf("Enter in the preview must not add: added=%+v", manager.added) + } +} + +func TestSameNamedServersMentionAndPromptCarryServerID(t *testing.T) { + imported := mcp.Server{Name: "docs", Scope: mcp.ScopeImported, Path: "/repo/.cursor/mcp.json"} + user := mcp.Server{Name: "docs", Scope: mcp.ScopeUser, Path: "/home/u/.rootly-cli/mcp.json"} + manager := &fakeMCPManager{views: []mcp.ServerView{ + {Server: imported, Status: mcp.StatusConnected}, + {Server: user, Status: mcp.StatusConnected, Prompts: []mcp.PromptInfo{{Name: "triage"}}, Resources: []mcp.ResourceInfo{{URI: "file:///runbook.md"}}}, + }} + model := NewTUIModel(context.Background(), &fakeStreamer{}, TUIOptions{NoColor: true, MCP: manager}) + model.mcpPage.views = manager.views + refs := model.resourceMentions("see @docs:file:///runbook.md") + if len(refs) != 1 || refs[0].serverID != user.ID() { + t.Fatalf("mention refs = %+v, want the user server's ID", refs) + } + if _, serverID, found := model.findPrompt("docs", "triage"); !found || serverID != user.ID() { + t.Fatalf("findPrompt serverID = %q found=%v", serverID, found) + } +} diff --git a/internal/agent/mcp_extras.go b/internal/agent/mcp_extras.go new file mode 100644 index 0000000..67905dc --- /dev/null +++ b/internal/agent/mcp_extras.go @@ -0,0 +1,361 @@ +package agent + +import ( + "context" + "fmt" + "regexp" + "strings" + "unicode" + + "charm.land/bubbles/v2/list" + tea "charm.land/bubbletea/v2" + + "github.com/rootlyhq/rootly-cli/internal/mcp" +) + +const ( + maxMentionSuggestions = 8 + maxPromptSuggestions = 8 + maxAttachedResources = 8 +) + +var resourceMentionPattern = regexp.MustCompile(`@([A-Za-z0-9_.-]+):(\S+)`) + +type mcpPromptMsg struct { + request int + label string + text string + err error +} + +type mentionsReadyMsg struct { + request int + echo string + message string + attached []string + notes []string +} + +type resourceRef struct { + server string + uri string + serverID string +} + +func (r resourceRef) mention() string { return "@" + r.server + ":" + r.uri } + +func (m *TUIModel) liveMCPViews() []mcp.ServerView { + if m.mcp == nil { + return nil + } + var views []mcp.ServerView + for _, view := range m.mcp.Views() { + if view.Status == mcp.StatusConnected && !view.SessionOff { + views = append(views, view) + } + } + return views +} + +func oneLine(text string) string { + return sanitizeTerminalText(strings.Join(strings.Fields(text), " ")) +} + +// appendPromptItems adds /server:prompt entries after built-ins and skills, which win on collision. +// paletteSafeName rejects prompt names that would print terminal controls or could not be typed back as one word. +func paletteSafeName(name string) bool { + return name != "" && name == sanitizeTerminalText(name) && !strings.ContainsFunc(name, unicode.IsSpace) +} + +func (m *TUIModel) appendPromptItems(items []list.Item, value string) []list.Item { + taken := map[string]bool{} + for _, item := range items { + if command, ok := item.(slashCommandItem); ok { + taken[strings.ToLower(command.command)] = true + } + } + lower := strings.ToLower(value) + added := 0 + for _, view := range m.liveMCPViews() { + for _, prompt := range view.Prompts { + if !paletteSafeName(prompt.Name) { + continue + } + command := "/" + view.Server.Name + ":" + prompt.Name + key := strings.ToLower(command) + if taken[key] || isBuiltinSlashCommand(command) || !strings.HasPrefix(key, lower) { + continue + } + description := oneLine(prompt.Description) + if description == "" { + description = "MCP prompt" + } + taken[key] = true + items = append(items, slashCommandItem{command: command, description: description, insert: len(prompt.Arguments) > 0}) + added++ + if added == maxPromptSuggestions { + return items + } + } + } + return items +} + +func (m *TUIModel) findPrompt(server, name string) (prompt mcp.PromptInfo, serverID string, found bool) { + for _, view := range m.liveMCPViews() { + if view.Server.Name != server { + continue + } + for _, candidate := range view.Prompts { + if candidate.Name == name { + return candidate, view.Server.ID(), true + } + } + } + return mcp.PromptInfo{}, "", false +} + +// promptArguments maps key=value tokens by name and fills the rest positionally; extra words join the last argument. +func promptArguments(prompt mcp.PromptInfo, tokens []string) (arguments map[string]string, missing []string) { + arguments = map[string]string{} + known := map[string]bool{} + for _, argument := range prompt.Arguments { + known[argument.Name] = true + } + var positional []string + for _, token := range tokens { + if key, value, ok := strings.Cut(token, "="); ok && known[key] { + arguments[key] = value + continue + } + positional = append(positional, token) + } + var open []string + for _, argument := range prompt.Arguments { + if _, set := arguments[argument.Name]; !set { + open = append(open, argument.Name) + } + } + for index, name := range open { + if index >= len(positional) { + break + } + if index == len(open)-1 { + arguments[name] = strings.Join(positional[index:], " ") + break + } + arguments[name] = positional[index] + } + for _, argument := range prompt.Arguments { + if argument.Required && arguments[argument.Name] == "" { + missing = append(missing, argument.Name) + } + } + return arguments, missing +} + +func promptUsage(command string, prompt mcp.PromptInfo) string { + parts := []string{command} + for _, argument := range prompt.Arguments { + if argument.Required { + parts = append(parts, "<"+argument.Name+">") + } else { + parts = append(parts, "["+argument.Name+"]") + } + } + return sanitizeTerminalText(strings.Join(parts, " ")) +} + +func (m *TUIModel) runMCPPrompt(command string) (tea.Cmd, bool) { + parts := strings.Fields(command) + if len(parts) == 0 || m.mcp == nil { + return nil, false + } + server, name, ok := strings.Cut(strings.TrimPrefix(parts[0], "/"), ":") + if !ok { + return nil, false + } + prompt, serverID, found := m.findPrompt(server, name) + if !found { + return nil, false + } + arguments, missing := promptArguments(prompt, parts[1:]) + if len(missing) > 0 { + return m.println("Usage: " + promptUsage(parts[0], prompt)), true + } + m.mentionRequest++ + manager, ctx, label, request := m.mcp, m.ctx, parts[0], m.mentionRequest + m.status = "Loading prompt " + sanitizeTerminalText(label) + "…" + return func() tea.Msg { + text, err := manager.GetPrompt(ctx, serverID, name, arguments) + return mcpPromptMsg{request: request, label: label, text: text, err: err} + }, true +} + +func (m *TUIModel) handleMCPPrompt(msg mcpPromptMsg) tea.Cmd { + if msg.request != m.mentionRequest { + return nil + } + m.status = "" + if msg.err != nil { + return m.println(faint(m.noColor, sanitizeTerminalText("Prompt "+msg.label+" failed: "+msg.err.Error()))) + } + return m.submitPrompt(sanitizeTerminalText(msg.text)) +} + +func (m *TUIModel) listMCPPrompts() tea.Cmd { + var lines []string + for _, view := range m.liveMCPViews() { + for _, prompt := range view.Prompts { + line := promptUsage("/"+view.Server.Name+":"+prompt.Name, prompt) + if description := oneLine(prompt.Description); description != "" { + line += " " + faint(m.noColor, description) + } + lines = append(lines, line) + } + } + if len(lines) == 0 { + return m.println("No MCP prompts from connected servers. Open /mcp to manage servers.") + } + return m.println("MCP prompts\n" + strings.Join(lines, "\n")) +} + +func (m *TUIModel) listMCPResources() tea.Cmd { + var lines []string + for _, view := range m.liveMCPViews() { + for _, resource := range view.Resources { + line := sanitizeTerminalText(resourceRef{server: view.Server.Name, uri: resource.URI}.mention()) + if name := oneLine(resource.Name); name != "" { + line += " " + faint(m.noColor, name) + } + lines = append(lines, line) + } + } + if len(lines) == 0 { + return m.println("No MCP resources from connected servers. Open /mcp to manage servers.") + } + return m.println("MCP resources (mention with @server:uri)\n" + strings.Join(lines, "\n")) +} + +func lastToken(value string) string { + if value == "" || strings.ContainsAny(value[len(value)-1:], " \t\r\n") { + return "" + } + return value[strings.LastIndexAny(value, " \t\r\n")+1:] +} + +// refreshMentionMenu opens the @server:uri palette for the word being typed. +func (m *TUIModel) refreshMentionMenu(value string) bool { + token := lastToken(value) + if !strings.HasPrefix(token, "@") || m.mcp == nil { + return false + } + query := strings.ToLower(token[1:]) + var prefixed, contains []list.Item + for _, view := range m.liveMCPViews() { + for _, resource := range view.Resources { + mention := resourceRef{server: view.Server.Name, uri: resource.URI}.mention() + item := slashCommandItem{command: sanitizeTerminalText(mention), description: oneLine(resource.Name), mention: true} + switch lower := strings.ToLower(mention[1:]); { + case strings.HasPrefix(lower, query): + prefixed = append(prefixed, item) + case strings.Contains(lower, query): + contains = append(contains, item) + } + } + } + items := make([]list.Item, 0, len(prefixed)+len(contains)) + items = append(items, prefixed...) + items = append(items, contains...) + if len(items) == 0 { + return false + } + m.showCommandMenu(items[:min(len(items), maxMentionSuggestions)]) + return true +} + +func (m *TUIModel) insertMention(mention string) { + value := m.composer.Value() + start := strings.LastIndexAny(value, " \t\r\n") + 1 + m.composer.SetValue(value[:start] + mention + " ") + m.composer.CursorEnd() +} + +func (m *TUIModel) resourceMentions(message string) []resourceRef { + known := map[resourceRef]string{} + for _, view := range m.liveMCPViews() { + for _, resource := range view.Resources { + ref := resourceRef{server: view.Server.Name, uri: resource.URI} + if _, taken := known[ref]; !taken { + known[ref] = view.Server.ID() + } + } + } + if len(known) == 0 { + return nil + } + var refs []resourceRef + seen := map[resourceRef]bool{} + for _, match := range resourceMentionPattern.FindAllStringSubmatch(message, -1) { + ref := resourceRef{server: match[1], uri: strings.TrimRight(match[2], ".,;:!?)")} + if _, ok := known[ref]; !ok { + ref.uri = match[2] + } + serverID, ok := known[ref] + ref.serverID = serverID + if ok && !seen[ref] && len(refs) < maxAttachedResources { + seen[ref] = true + refs = append(refs, ref) + } + } + return refs +} + +// sendWithResources reads mentioned resources before sending; nil when the message mentions none. +func (m *TUIModel) sendWithResources(message string) tea.Cmd { + refs := m.resourceMentions(message) + if len(refs) == 0 { + return nil + } + m.mentionRequest++ + request, manager, ctx := m.mentionRequest, m.mcp, m.ctx + m.status = fmt.Sprintf("Reading %d MCP resource(s)…", len(refs)) + return func() tea.Msg { + return readMentionedResources(ctx, manager, request, message, refs) + } +} + +func readMentionedResources(ctx context.Context, manager MCPManager, request int, message string, refs []resourceRef) mentionsReadyMsg { + result := mentionsReadyMsg{request: request, echo: message, message: message} + var blocks []string + for _, ref := range refs { + text, truncated, err := manager.ReadResource(ctx, ref.serverID, ref.uri, mcp.ResourceByteCap) + if err != nil { + result.notes = append(result.notes, "skipped "+ref.mention()+": "+err.Error()) + continue + } + label := ref.mention() + if truncated { + label += fmt.Sprintf(" (first %d KB)", mcp.ResourceByteCap/1024) + } + result.attached = append(result.attached, label) + blocks = append(blocks, fmt.Sprintf("\n%s\n", ref.server, ref.uri, text)) + } + if len(blocks) > 0 { + result.message = strings.Join(blocks, "\n\n") + "\n\n" + message + } + return result +} + +func (m *TUIModel) handleMentionsReady(msg mentionsReadyMsg) tea.Cmd { + if msg.request != m.mentionRequest { + return nil + } + m.status = "" + var notes []string + if len(msg.attached) > 0 { + notes = append(notes, "attached: "+strings.Join(msg.attached, ", ")) + } + notes = append(notes, msg.notes...) + return m.sendPrompt(msg.echo, msg.message, sanitizeTerminalText(strings.Join(notes, "\n"))) +} diff --git a/internal/agent/mcp_extras_test.go b/internal/agent/mcp_extras_test.go new file mode 100644 index 0000000..2ca93ba --- /dev/null +++ b/internal/agent/mcp_extras_test.go @@ -0,0 +1,293 @@ +package agent + +import ( + "context" + "fmt" + "strings" + "testing" + + "charm.land/bubbles/v2/list" + tea "charm.land/bubbletea/v2" + "github.com/google/uuid" + rootly "github.com/rootlyhq/rootly-go" + + "github.com/rootlyhq/rootly-cli/internal/mcp" +) + +func extrasManager() *fakeMCPManager { + resources := []mcp.ResourceInfo{{URI: "file:///runbook.md", Name: "Runbook"}, {URI: "file:///big.log", Name: "Big log"}, {URI: "file:///gone.md", Name: "Gone"}} + for index := range 10 { + resources = append(resources, mcp.ResourceInfo{URI: fmt.Sprintf("file:///extra-%d.md", index), Name: "Extra"}) + } + return &fakeMCPManager{ + views: []mcp.ServerView{ + { + Server: mcp.Server{Name: "docs", Scope: mcp.ScopeUser}, Status: mcp.StatusConnected, + Tools: []mcp.ToolInfo{{Name: "search"}}, + Prompts: []mcp.PromptInfo{{Name: "triage", Description: "Triage an\nalert", Arguments: []mcp.PromptArg{{Name: "service", Required: true}, {Name: "notes"}}}, {Name: "summary"}}, + Resources: resources, + }, + {Server: mcp.Server{Name: "offline", Scope: mcp.ScopeUser}, Status: mcp.StatusError, Prompts: []mcp.PromptInfo{{Name: "hidden"}}}, + }, + resources: map[string]string{ + "docs:file:///runbook.md": "restart the api", + "docs:file:///big.log": strings.Repeat("x", mcp.ResourceByteCap+500), + }, + } +} + +func menuCommands(model *TUIModel) []string { + var commands []string + for _, item := range model.commandMenu.Items() { + commands = append(commands, item.(slashCommandItem).command) + } + return commands +} + +func TestSlashPaletteListsMCPPromptsFromConnectedServers(t *testing.T) { + model := newMCPTestModel(t, extrasManager(), &fakeConnectorAPI{}) + model.composer.SetValue("/") + model.refreshCommandMenu() + model.composer.SetValue("/docs") + model.refreshCommandMenu() + got := menuCommands(model) + if !model.commandMenuOpen || strings.Join(got, ",") != "/docs:triage,/docs:summary" { + t.Fatalf("menu = %v (open %v)", got, model.commandMenuOpen) + } + model.composer.SetValue("/off") + model.refreshCommandMenu() + if model.commandMenuOpen { + t.Fatalf("prompts from disconnected servers listed: %v", menuCommands(model)) + } +} + +func TestPromptItemsLoseToBuiltinsAndSkills(t *testing.T) { + model := newMCPTestModel(t, extrasManager(), &fakeConnectorAPI{}) + items := model.appendPromptItems(nil, "/") + if len(items) != 2 { + t.Fatalf("items = %v", items) + } + skill := slashCommandItem{command: "/docs:triage", description: "skill", skill: true} + items = model.appendPromptItems([]list.Item{skill}, "/docs:t") + if len(items) != 1 || items[0].(slashCommandItem) != skill { + t.Fatalf("skill should win on collision, got %v", items) + } + if !isBuiltinSlashCommand("/prompts") || !isBuiltinSlashCommand("/resources") { + t.Fatal("/prompts and /resources must be built-ins so skills cannot shadow them") + } +} + +func TestRunMCPPromptArgumentsAndUsage(t *testing.T) { + model := newMCPTestModel(t, extrasManager(), &fakeConnectorAPI{}) + cmd := model.runSlashCommand("/docs:triage") + runCmd(model, cmd) + if !strings.Contains(model.Transcript(), "Usage: /docs:triage [notes]") { + t.Fatalf("transcript = %s", model.Transcript()) + } + cmd = model.runSlashCommand("/docs:triage api notes=urgent") + msg, ok := cmd().(mcpPromptMsg) + if !ok || msg.err != nil || msg.text != "prompt triage map[notes:urgent service:api]" { + t.Fatalf("prompt msg = %+v", msg) + } + model.Update(msg) + if !strings.Contains(model.Transcript(), "prompt triage map[notes:urgent service:api]") { + t.Fatalf("prompt not sent: %s", model.Transcript()) + } + if _, missing := promptArguments(mcp.PromptInfo{Arguments: []mcp.PromptArg{{Name: "service", Required: true}, {Name: "notes"}}}, []string{"api", "very", "slow"}); missing != nil { + t.Fatalf("missing = %v", missing) + } + args, _ := promptArguments(mcp.PromptInfo{Arguments: []mcp.PromptArg{{Name: "service"}, {Name: "notes"}}}, []string{"api", "very", "slow"}) + if args["service"] != "api" || args["notes"] != "very slow" { + t.Fatalf("args = %v", args) + } +} + +func TestPromptsAndResourcesListings(t *testing.T) { + model := newMCPTestModel(t, extrasManager(), &fakeConnectorAPI{}) + model.runSlashCommand("/prompts") + model.runSlashCommand("/resources") + transcript := model.Transcript() + for _, want := range []string{"/docs:triage [notes] Triage an alert", "/docs:summary", "@docs:file:///runbook.md Runbook"} { + if !strings.Contains(transcript, want) { + t.Fatalf("missing %q in:\n%s", want, transcript) + } + } + if strings.Contains(transcript, "hidden") { + t.Fatalf("listed prompt from a disconnected server:\n%s", transcript) + } + empty := newMCPTestModel(t, &fakeMCPManager{}, &fakeConnectorAPI{}) + empty.runSlashCommand("/prompts") + if !strings.Contains(empty.Transcript(), "No MCP prompts") { + t.Fatalf("transcript = %s", empty.Transcript()) + } +} + +func TestMentionPaletteShowsTopEightAndTabInserts(t *testing.T) { + model := newMCPTestModel(t, extrasManager(), &fakeConnectorAPI{}) + model.composer.SetValue("check @docs") + model.refreshCommandMenu() + if got := menuCommands(model); !model.commandMenuOpen || len(got) != maxMentionSuggestions { + t.Fatalf("menu = %v", got) + } + model.composer.SetValue("check @run") + model.refreshCommandMenu() + if got := menuCommands(model); len(got) != 1 || got[0] != "@docs:file:///runbook.md" { + t.Fatalf("menu = %v", got) + } + model.handleCommandMenuKey(tea.KeyPressMsg{Code: tea.KeyTab}) + if got := model.composer.Value(); got != "check @docs:file:///runbook.md " { + t.Fatalf("composer = %q", got) + } + model.composer.SetValue("mail me@") + model.refreshCommandMenu() + if model.commandMenuOpen && len(menuCommands(model)) > maxMentionSuggestions { + t.Fatal("menu too long") + } +} + +func TestMentionedResourcesAttachedWithCapAndNotes(t *testing.T) { + manager := extrasManager() + model := newMCPTestModel(t, manager, &fakeConnectorAPI{}) + if model.sendWithResources("email ops@example.com and @docs:unknown") != nil { + t.Fatal("unknown mentions should send as plain text") + } + cmd := model.sendWithResources("why? @docs:file:///runbook.md, @docs:file:///big.log @docs:file:///gone.md @docs:file:///runbook.md") + msg, ok := cmd().(mentionsReadyMsg) + if !ok { + t.Fatal("expected mentionsReadyMsg") + } + if strings.Join(msg.attached, "|") != "@docs:file:///runbook.md|@docs:file:///big.log (first 32 KB)" { + t.Fatalf("attached = %v", msg.attached) + } + if len(msg.notes) != 1 || !strings.Contains(msg.notes[0], "skipped @docs:file:///gone.md: no text content") { + t.Fatalf("notes = %v", msg.notes) + } + if !strings.HasPrefix(msg.message, "\nrestart the api\n") || + !strings.HasSuffix(msg.message, "\n\n"+msg.echo) || strings.Count(msg.message, " mcp.ResourceByteCap+10 { + t.Fatal("resource cap not applied") + } + for _, limit := range manager.readLimits { + if limit != mcp.ResourceByteCap { + t.Fatalf("read limit = %d", limit) + } + } + model.Update(msg) + transcript := model.Transcript() + if !strings.Contains(transcript, "attached: @docs:file:///runbook.md, @docs:file:///big.log (first 32 KB)") || !strings.Contains(transcript, "skipped @docs:file:///gone.md") { + t.Fatalf("transcript = %s", transcript) + } + if strings.Contains(transcript, "restart the api") { + t.Fatal("resource body should not be echoed into the transcript") + } +} + +func TestMCPRowCounts(t *testing.T) { + view := mcp.ServerView{Status: mcp.StatusConnected, Tools: make([]mcp.ToolInfo, 3), Resources: make([]mcp.ResourceInfo, 1), Prompts: make([]mcp.PromptInfo, 2)} + if got := mcpStatusLabel(view); got != "✔ 3 tools · 1 resource · 2 prompts" { + t.Fatalf("label = %q", got) + } + if got := mcpStatusLabel(mcp.ServerView{Status: mcp.StatusConnected, Tools: make([]mcp.ToolInfo, 1)}); got != "✔ 1 tool" { + t.Fatalf("label = %q", got) + } +} + +func TestStalePromptAndMentionLoadsAreDropped(t *testing.T) { + model := newMCPTestModel(t, extrasManager(), &fakeConnectorAPI{}) + pending := model.runSlashCommand("/docs:triage api") + model.startNewSession() + model.Update(pending()) + if strings.Contains(model.Transcript(), "prompt triage") { + t.Fatalf("prompt loaded before /new was sent into the new chat:\n%s", model.Transcript()) + } + pending = model.runSlashCommand("/docs:triage api") + model.sendPrompt("other question", "other question", "") + model.Update(pending()) + if strings.Contains(model.Transcript(), "prompt triage") { + t.Fatalf("prompt loaded before another send replaced it:\n%s", model.Transcript()) + } +} + +func TestSessionSwitchDropsPendingPromptAndMentionLoads(t *testing.T) { + model := newMCPTestModel(t, extrasManager(), &fakeConnectorAPI{}) + pending := model.runSlashCommand("/docs:triage api") + model.selectPickerItem(sessionPickerItem{session: testSession(uuid.NewString(), "Other chat")}) + model.historyPending = false + model.Update(pending()) + if model.state == tuiStreaming || strings.Contains(model.Transcript(), "prompt triage") { + t.Fatal("a prompt loaded before resuming another session was sent into it") + } + model.composer.SetValue("see @docs:file:///runbook.md") + pending = model.submit() + model.Update(bindingResolvedMsg{request: model.actionRequest}) + model.Update(pending()) + if model.state == tuiStreaming { + t.Fatal("a resource load started before switching the bound incident was still sent") + } +} + +func TestCtrlCCancelsPendingPromptAndMentionLoads(t *testing.T) { + model := newMCPTestModel(t, extrasManager(), &fakeConnectorAPI{}) + pending := model.runSlashCommand("/docs:triage api") + model.Update(tea.KeyPressMsg{Code: 'c', Mod: tea.ModCtrl}) + model.Update(pending()) + if model.state == tuiStreaming || strings.Contains(model.Transcript(), "prompt triage") { + t.Fatal("a prompt load cancelled with Ctrl+C was still sent") + } + model.composer.SetValue("see @docs:file:///runbook.md") + pending = model.submit() + model.Update(tea.KeyPressMsg{Code: 'c', Mod: tea.ModCtrl}) + model.Update(pending()) + if model.state == tuiStreaming { + t.Fatal("a resource load cancelled with Ctrl+C was still sent") + } +} + +func TestPromptNamesWithTerminalControlsStayOutOfThePalette(t *testing.T) { + manager := extrasManager() + manager.views[0].Prompts = append(manager.views[0].Prompts, mcp.PromptInfo{Name: "evil\x1b]0;pwned\x07"}, mcp.PromptInfo{Name: "two words"}) + model := newMCPTestModel(t, manager, &fakeConnectorAPI{}) + for _, item := range model.appendPromptItems(nil, "/") { + command := item.(slashCommandItem).command + if strings.Contains(command, "\x1b") || strings.Contains(command, "two words") { + t.Fatalf("unsafe prompt name reached the palette: %q", command) + } + } + if len(model.appendPromptItems(nil, "/docs:")) != 2 { + t.Fatal("safe prompts should still be listed") + } +} + +func TestMCPListRowShowsFullCountsAtTypicalWidth(t *testing.T) { + manager := &fakeMCPManager{views: []mcp.ServerView{{ + Server: mcp.Server{Name: "everything", Scope: mcp.ScopeUser}, Status: mcp.StatusConnected, + Tools: make([]mcp.ToolInfo, 22), Resources: make([]mcp.ResourceInfo, 11), Prompts: make([]mcp.PromptInfo, 4), + }}} + model := newMCPTestModel(t, manager, &fakeConnectorAPI{}) + model.width = 139 + model.mcpPage.open = true + if view := model.mcpPageView(); !strings.Contains(view, "✔ 22 tools · 11 resources · 4 prompts") { + t.Fatalf("list row clipped the capability counts:\n%s", view) + } +} + +func TestMCPPromptTextLosesTerminalControls(t *testing.T) { + manager := extrasManager() + manager.promptText = "clear\x1b[2J\x1b[H then\x1b]52;c;cGF5bG9hZA==\x07 copy \x1b]0;title\x1b\\ and \x9b31m forge" + model := NewTUIModel(context.Background(), &fakeStreamer{responses: [][]rootly.AiChatStreamEvent{{}}}, TUIOptions{NoColor: true, API: &fakeConnectorAPI{}, MCP: manager}) + model.Update(tea.WindowSizeMsg{Width: 120, Height: 40}) + cmd := model.handleMCPPrompt(mcpPromptMsg{request: model.mentionRequest, label: "/docs:summarize", text: manager.promptText}) + if cmd == nil || model.state != tuiStreaming { + t.Fatal("prompt was not submitted") + } + for _, text := range []string{transcriptText(model), model.currentPrompt, model.promptHistory[len(model.promptHistory)-1], *model.currentParams.Message} { + if strings.ContainsAny(text, "\x1b\x07") || strings.Contains(text, "\x9b") || strings.Contains(text, "]52;") { + t.Fatalf("terminal control survived: %q", text) + } + } + if !strings.Contains(transcriptText(model), "clear then copy and 31m forge") { + t.Fatalf("prompt text was lost: %q", transcriptText(model)) + } +} diff --git a/internal/agent/mcp_page.go b/internal/agent/mcp_page.go new file mode 100644 index 0000000..525ab86 --- /dev/null +++ b/internal/agent/mcp_page.go @@ -0,0 +1,1066 @@ +package agent + +import ( + "context" + "fmt" + "os" + "strconv" + "strings" + + tea "charm.land/bubbletea/v2" + "charm.land/lipgloss/v2" + "github.com/charmbracelet/x/ansi" + + "github.com/rootlyhq/rootly-cli/internal/api" + "github.com/rootlyhq/rootly-cli/internal/mcp" +) + +const cloudConnectorPageSize = 25 + +// MCPManager is the agent's handle on local MCP servers. +type MCPManager interface { + LocalTools + Views() []mcp.ServerView + BuiltinTools() []mcp.BuiltinToolInfo + SetBuiltinEnabled(string, bool) error + Problems() []string + Changes() <-chan struct{} + Connecting() int + ToggleSession(string) error + GetPrompt(ctx context.Context, serverID, name string, arguments map[string]string) (string, error) + ReadResource(ctx context.Context, serverID, uri string, limit int) (string, bool, error) + Approve(string) error + Enable(string) error + Reconnect(string) error + Login(context.Context, string, mcp.LoginOptions) error + AddFromMarket(mcp.MarketEntry) (mcp.Server, error) +} + +type ConnectorLister interface { + ListAIConnectors(context.Context, int, int) (*api.AIConnectorPage, error) +} + +type mcpChangedMsg struct{} + +type cloudConnectorsMsg struct { + request int + page *api.AIConnectorPage + err error +} + +type mcpLoginMsg struct { + id string + url string + err error + done bool + events <-chan mcpLoginMsg +} + +type mcpRowKind int + +const ( + mcpRowHeader mcpRowKind = iota + mcpRowServer + mcpRowCloud + mcpRowLoadMore + mcpRowNote +) + +type mcpRow struct { + kind mcpRowKind + text string + view *mcp.ServerView + cloud *api.AIConnector +} + +type mcpPage struct { + open bool + cursor int + offset int + filtering bool + filter string + detail string + detailCursor int + detailOffset int + detailFiltering bool + detailFilter string + cloud []api.AIConnector + cloudNext int + cloudLoading bool + cloudLoaded bool + cloudErr string + cloudRequest int + message string + loginURL string + views []mcp.ServerView + market *marketState +} + +func (m *TUIModel) watchMCP() tea.Cmd { + if m.mcp == nil { + return nil + } + changes := m.mcp.Changes() + return func() tea.Msg { + if _, ok := <-changes; !ok { + return nil + } + return mcpChangedMsg{} + } +} + +func (m *TUIModel) openMCPPage() tea.Cmd { + m.mcpPage.open = true + m.mcpPage.cursor, m.mcpPage.offset = 0, 0 + m.mcpPage.filtering, m.mcpPage.filter = false, "" + m.mcpPage.detail = "" + m.mcpPage.message = "" + m.refreshMCPViews() + if !m.mcpPage.cloudLoading { + m.mcpPage.cloud, m.mcpPage.cloudLoaded = nil, false + return m.loadCloudConnectors(1) + } + return nil +} + +func (m *TUIModel) refreshMCPViews() { + if m.mcp == nil { + m.mcpPage.views = nil + return + } + m.mcpPage.views = m.mcp.Views() + builtinTools := m.mcp.BuiltinTools() + if len(builtinTools) == 0 { + return + } + builtinView := mcp.ServerView{ + Server: mcp.Server{Name: "Built-in", Scope: mcp.ScopeBuiltin, Path: "builtin"}, + Status: mcp.StatusConnected, + Tools: make([]mcp.ToolInfo, 0, len(builtinTools)), + } + for _, tool := range builtinTools { + builtinView.Tools = append(builtinView.Tools, mcp.ToolInfo(tool)) + } + m.mcpPage.views = append([]mcp.ServerView{builtinView}, m.mcpPage.views...) +} + +func (m *TUIModel) loadCloudConnectors(page int) tea.Cmd { + lister, ok := m.api.(ConnectorLister) + if !ok { + m.mcpPage.cloudLoaded = true + m.mcpPage.cloudErr = "unavailable" + return nil + } + m.mcpPage.cloudLoading = true + m.mcpPage.cloudRequest++ + request := m.mcpPage.cloudRequest + ctx := m.ctx + return func() tea.Msg { + result, err := lister.ListAIConnectors(ctx, page, cloudConnectorPageSize) + return cloudConnectorsMsg{request: request, page: result, err: err} + } +} + +func (m *TUIModel) handleCloudConnectors(msg cloudConnectorsMsg) { + if msg.request != m.mcpPage.cloudRequest { + return + } + m.mcpPage.cloudLoading = false + m.mcpPage.cloudLoaded = true + if msg.err != nil { + m.mcpPage.cloudErr = sanitizeTerminalText(msg.err.Error()) + if msg.page != nil && (msg.page.StatusCode == 403 || msg.page.StatusCode == 404) { + m.mcpPage.cloudErr = "not available for your account" + } + return + } + m.mcpPage.cloudErr = "" + m.mcpPage.cloud = append(m.mcpPage.cloud, msg.page.Connectors...) + m.mcpPage.cloudNext = msg.page.NextPage +} + +func (m *TUIModel) mcpRows() []mcpRow { + page := &m.mcpPage + filter := strings.ToLower(page.filter) + match := func(name string) bool { return filter == "" || strings.Contains(strings.ToLower(name), filter) } + rows := []mcpRow{{kind: mcpRowHeader, text: "Rootly cloud"}} + cloudRows := 0 + for index := range page.cloud { + if match(page.cloud[index].Name) { + rows = append(rows, mcpRow{kind: mcpRowCloud, cloud: &page.cloud[index]}) + cloudRows++ + } + } + switch { + case page.cloudLoading: + rows = append(rows, mcpRow{kind: mcpRowNote, text: "Loading…"}) + case page.cloudErr != "": + rows = append(rows, mcpRow{kind: mcpRowNote, text: "Could not load cloud connectors: " + page.cloudErr}) + case page.cloudNext > 0: + rows = append(rows, mcpRow{kind: mcpRowLoadMore, text: "Load more…"}) + case cloudRows == 0 && page.cloudLoaded: + rows = append(rows, mcpRow{kind: mcpRowNote, text: "No cloud connectors"}) + } + groups := []struct { + title string + scopes []string + }{ + {"Built-in", []string{mcp.ScopeBuiltin}}, + {"Project MCPs", []string{mcp.ScopeProject, mcp.ScopeLocal}}, + {"User MCPs (~/.rootly-cli/mcp.json)", []string{mcp.ScopeUser}}, + {"Imported (read-only)", []string{mcp.ScopeImported}}, + } + for _, group := range groups { + title := group.title + if group.scopes[0] == mcp.ScopeProject { + if dir, err := os.Getwd(); err == nil { + title += " (" + displayHomePath(dir) + ")" + } + } + groupRows := []mcpRow{} + for index := range page.views { + view := &page.views[index] + if !containsString(group.scopes, view.Server.Scope) || !matchesMCPView(view, match) { + continue + } + groupRows = append(groupRows, mcpRow{kind: mcpRowServer, view: view}) + } + if len(groupRows) == 0 && filter != "" { + continue + } + rows = append(rows, mcpRow{kind: mcpRowHeader, text: title}) + if len(groupRows) == 0 { + hint := "None · rootly mcp add …" + switch group.scopes[0] { + case mcp.ScopeImported: + hint = "None found in Claude Code, Codex, Cursor, VS Code or Claude Desktop" + case mcp.ScopeBuiltin: + hint = "No built-in tools available" + } + rows = append(rows, mcpRow{kind: mcpRowNote, text: hint}) + } + rows = append(rows, groupRows...) + } + return rows +} + +func matchesMCPView(view *mcp.ServerView, match func(string) bool) bool { + if match(view.Server.Name) { + return true + } + if view.Server.Scope != mcp.ScopeBuiltin { + return false + } + for _, tool := range view.Tools { + if match(tool.Name) { + return true + } + } + return false +} + +func containsString(values []string, value string) bool { + for _, candidate := range values { + if candidate == value { + return true + } + } + return false +} + +func displayHomePath(path string) string { + home, err := os.UserHomeDir() + if err == nil && home != "" && (path == home || strings.HasPrefix(path, home+string(os.PathSeparator))) { + return "~" + strings.TrimPrefix(path, home) + } + return path +} + +func selectable(row mcpRow) bool { + return row.kind == mcpRowServer || row.kind == mcpRowCloud || row.kind == mcpRowLoadMore +} + +func (m *TUIModel) moveMCPCursor(rows []mcpRow, delta int) { + if len(rows) == 0 { + return + } + index := m.mcpPage.cursor + for step := 0; step < len(rows); step++ { + index += delta + if index < 0 || index >= len(rows) { + return + } + if selectable(rows[index]) { + m.mcpPage.cursor = index + return + } + } +} + +func (m *TUIModel) clampMCPCursor(rows []mcpRow) { + if m.mcpPage.cursor >= len(rows) { + m.mcpPage.cursor = len(rows) - 1 + } + if m.mcpPage.cursor < 0 { + m.mcpPage.cursor = 0 + } + if len(rows) > 0 && !selectable(rows[m.mcpPage.cursor]) { + for index := m.mcpPage.cursor; index < len(rows); index++ { + if selectable(rows[index]) { + m.mcpPage.cursor = index + return + } + } + for index := m.mcpPage.cursor; index >= 0; index-- { + if selectable(rows[index]) { + m.mcpPage.cursor = index + return + } + } + } +} + +func (m *TUIModel) selectedMCPRow() (mcpRow, bool) { + rows := m.mcpRows() + m.clampMCPCursor(rows) + if len(rows) == 0 || !selectable(rows[m.mcpPage.cursor]) { + return mcpRow{}, false + } + return rows[m.mcpPage.cursor], true +} + +func (m *TUIModel) detailView() (*mcp.ServerView, *api.AIConnector) { + id := m.mcpPage.detail + if strings.HasPrefix(id, "cloud:") { + for index := range m.mcpPage.cloud { + if "cloud:"+m.mcpPage.cloud[index].ID == id { + return nil, &m.mcpPage.cloud[index] + } + } + return nil, nil + } + for index := range m.mcpPage.views { + if m.mcpPage.views[index].Server.ID() == id { + return &m.mcpPage.views[index], nil + } + } + return nil, nil +} + +func (m *TUIModel) handleMCPPageKey(msg tea.KeyPressMsg) (tea.Model, tea.Cmd) { + page := &m.mcpPage + key := msg.String() + if page.detail != "" { + return m.handleMCPDetailKey(msg) + } + if page.market != nil { + cmd := m.handleMarketKey(msg) + return m, cmd + } + if page.filtering { + switch key { + case "esc": + page.filtering, page.filter = false, "" + case "enter": + page.filtering = false + case "backspace": + if page.filter != "" { + page.filter = string([]rune(page.filter)[:len([]rune(page.filter))-1]) + } + default: + if msg.Text != "" { + page.filter += msg.Text + } + } + page.cursor, page.offset = 0, 0 + m.clampMCPCursor(m.mcpRows()) + return m, nil + } + rows := m.mcpRows() + m.clampMCPCursor(rows) + switch key { + case "esc", "q", "ctrl+c": + page.open = false + return m, nil + case "up", "k": + m.moveMCPCursor(rows, -1) + case "down", "j": + m.moveMCPCursor(rows, 1) + case "pgup": + for range max(1, m.mcpListHeight()-1) { + m.moveMCPCursor(rows, -1) + } + case "pgdown": + for range max(1, m.mcpListHeight()-1) { + m.moveMCPCursor(rows, 1) + } + case "home", "g": + page.cursor = 0 + m.clampMCPCursor(rows) + case "end", "G": + page.cursor = len(rows) - 1 + m.clampMCPCursor(rows) + case "/": + page.filtering = true + case "m": + cmd := m.openMarket() + return m, cmd + case "enter": + row, ok := m.selectedMCPRow() + if !ok { + return m, nil + } + switch row.kind { + case mcpRowLoadMore: + cmd := m.loadCloudConnectors(page.cloudNext) + return m, cmd + case mcpRowCloud: + page.detail = "cloud:" + row.cloud.ID + case mcpRowServer: + page.detail = row.view.Server.ID() + if row.view.Status == mcp.StatusNeedsAuth { + page.message = "Needs authentication: press l to log in" + } + } + page.detailCursor, page.detailOffset, page.detailFilter, page.detailFiltering = 0, 0, "", false + case "space", " ", "a", "l", "r": + row, ok := m.selectedMCPRow() + if ok && row.kind == mcpRowServer { + if row.view.Server.Scope == mcp.ScopeBuiltin && (key == "space" || key == " ") { + return m, nil + } + cmd := m.mcpAction(row.view, key) + return m, cmd + } + } + return m, nil +} + +func (m *TUIModel) handleMCPDetailKey(msg tea.KeyPressMsg) (tea.Model, tea.Cmd) { + page := &m.mcpPage + key := msg.String() + view, _ := m.detailView() + if page.detailFiltering { + switch key { + case "esc": + page.detailFiltering, page.detailFilter = false, "" + case "enter": + page.detailFiltering = false + case "backspace": + if page.detailFilter != "" { + page.detailFilter = string([]rune(page.detailFilter)[:len([]rune(page.detailFilter))-1]) + } + default: + if msg.Text != "" { + page.detailFilter += msg.Text + } + } + page.detailCursor, page.detailOffset = 0, 0 + return m, nil + } + tools := m.detailTools(view) + switch key { + case "esc", "q", "backspace": + page.detail, page.message = "", "" + case "ctrl+c": + page.detail, page.open = "", false + case "up", "k": + page.detailCursor = max(0, page.detailCursor-1) + case "down", "j": + page.detailCursor = min(max(0, len(tools)-1), page.detailCursor+1) + case "pgup": + page.detailCursor = max(0, page.detailCursor-m.mcpListHeight()) + case "pgdown": + page.detailCursor = min(max(0, len(tools)-1), page.detailCursor+m.mcpListHeight()) + case "/": + page.detailFiltering = true + case "space", " ", "a", "l", "r": + if view != nil { + cmd := m.mcpAction(view, key) + return m, cmd + } + } + return m, nil +} + +func (m *TUIModel) mcpAction(view *mcp.ServerView, key string) tea.Cmd { + if m.mcp == nil { + return nil + } + page := &m.mcpPage + if view.Server.Scope == mcp.ScopeBuiltin { + if key == "space" || key == " " { + m.toggleBuiltinTool(view) + } + return nil + } + id := view.Server.ID() + name := sanitizeTerminalText(view.Server.Name) + var err error + switch key { + case "space", " ": + if view.Status == mcp.StatusOff { + err = m.mcp.Enable(id) + page.message = "Enabled " + name + } else { + err = m.mcp.ToggleSession(id) + if view.SessionOff { + page.message = name + " is on for this session" + } else { + page.message = name + " is off for this session" + } + } + case "a": + err = m.mcp.Approve(id) + page.message = "Approved " + name + case "r": + err = m.mcp.Reconnect(id) + page.message = "Reconnecting " + name + "…" + case "l": + if view.Server.Config.Transport() == mcp.TransportStdio { + page.message = name + " is a stdio server; there is nothing to log in to" + return nil + } + page.message = "Opening your browser to log in to " + name + "…" + return m.startMCPLogin(id) + } + if err != nil { + page.message = sanitizeTerminalText(err.Error()) + } + m.refreshMCPViews() + return nil +} + +func (m *TUIModel) toggleBuiltinTool(view *mcp.ServerView) { + tools := m.detailTools(view) + if len(tools) == 0 || m.mcpPage.detailCursor >= len(tools) { + return + } + tool := tools[m.mcpPage.detailCursor] + enabled := !tool.Enabled + if err := m.mcp.SetBuiltinEnabled(tool.Name, enabled); err != nil { + m.mcpPage.message = sanitizeTerminalText(err.Error()) + } else if enabled { + m.mcpPage.message = "Enabled " + sanitizeTerminalText(tool.Name) + } else { + m.mcpPage.message = "Disabled " + sanitizeTerminalText(tool.Name) + } + m.refreshMCPViews() +} + +func (m *TUIModel) startMCPLogin(id string) tea.Cmd { + manager := m.mcp + ctx := m.ctx + events := make(chan mcpLoginMsg, 2) + go func() { + defer close(events) + err := manager.Login(ctx, id, mcp.LoginOptions{OpenURL: func(url string) { + _ = mcp.OpenBrowser(url) + events <- mcpLoginMsg{id: id, url: url} + }}) + events <- mcpLoginMsg{id: id, err: err, done: true} + }() + return waitMCPLogin(events) +} + +func waitMCPLogin(events <-chan mcpLoginMsg) tea.Cmd { + return func() tea.Msg { + msg, ok := <-events + if !ok { + return nil + } + msg.events = events + return msg + } +} + +func (m *TUIModel) handleMCPLogin(msg mcpLoginMsg) tea.Cmd { + if !msg.done { + m.mcpPage.loginURL = msg.url + return waitMCPLogin(msg.events) + } + m.mcpPage.loginURL = "" + if msg.err != nil { + m.mcpPage.message = "Login failed: " + sanitizeTerminalText(msg.err.Error()) + } else { + m.mcpPage.message = "Logged in. Reconnecting…" + } + m.refreshMCPViews() + return nil +} + +func (m *TUIModel) detailTools(view *mcp.ServerView) []mcp.ToolInfo { + if view == nil { + return nil + } + filter := strings.ToLower(m.mcpPage.detailFilter) + if filter == "" { + return view.Tools + } + tools := make([]mcp.ToolInfo, 0, len(view.Tools)) + for _, tool := range view.Tools { + if strings.Contains(strings.ToLower(tool.Name), filter) || strings.Contains(strings.ToLower(tool.Description), filter) { + tools = append(tools, tool) + } + } + return tools +} + +func (m *TUIModel) mcpListHeight() int { + return max(3, m.height-6) +} + +func mcpStatusLabel(view mcp.ServerView) string { + if view.Server.Scope == mcp.ScopeBuiltin { + enabled := 0 + for _, tool := range view.Tools { + if tool.Enabled { + enabled++ + } + } + return fmt.Sprintf("%d/%d enabled", enabled, len(view.Tools)) + } + if view.SessionOff { + return "○ off for this session" + } + switch view.Status { + case mcp.StatusConnected: + return "✔ " + mcpCounts(view) + case mcp.StatusError: + return "✘ error" + case mcp.StatusNeedsAuth: + if view.Detail != "" && strings.Contains(view.Detail, "reconnect") { + return "▲ reconnect" + } + return "▲ needs authentication" + case mcp.StatusPending: + return "‖ pending approval" + case mcp.StatusConnecting: + return "… connecting" + default: + return "○ off" + } +} + +func cloudStatusLabel(connector api.AIConnector) string { + if !connector.AvailableToAgent { + status := connector.Status + if status == "" { + status = "off" + } + return "○ " + sanitizeTerminalText(status) + } + if connector.ToolCount != nil { + return "✔ " + pluralTools(*connector.ToolCount) + } + return "✔ available" +} + +func pluralTools(count int) string { + if count == 1 { + return "1 tool" + } + return strconv.Itoa(count) + " tools" +} + +func (m *TUIModel) mcpSummary() string { + servers, tools := len(m.mcpPage.cloud), 0 + for _, connector := range m.mcpPage.cloud { + if connector.AvailableToAgent && connector.ToolCount != nil { + tools += *connector.ToolCount + } + } + for _, view := range m.mcpPage.views { + if view.Server.Scope == mcp.ScopeBuiltin { + for _, tool := range view.Tools { + if tool.Enabled { + tools++ + } + } + continue + } + servers++ + if view.Status == mcp.StatusConnected && !view.SessionOff { + tools += len(view.Tools) + } + } + return fmt.Sprintf("%d servers · %d tools available to the agent", servers, tools) +} + +func (m *TUIModel) mcpStyle(style lipgloss.Style, text string) string { + if m.noColor { + return text + } + return style.Render(text) +} + +func (m *TUIModel) mcpPageView() string { + if m.mcpPage.detail != "" { + return m.mcpDetailView() + } + if m.mcpPage.market != nil { + return m.marketView() + } + width := max(20, m.width-1) + bold := lipgloss.NewStyle().Bold(true) + lines := []string{m.mcpStyle(bold, "Manage MCP servers"), faint(m.noColor, m.mcpSummary()), ""} + if m.mcp != nil { + if problems := m.mcp.Problems(); len(problems) > 0 { + lines[1] += faint(m.noColor, " · ✘ config: "+sanitizeTerminalText(strings.Join(problems, "; "))) + } + } + if m.mcpPage.filtering || m.mcpPage.filter != "" { + lines[2] = "Filter: " + sanitizeTerminalText(m.mcpPage.filter) + map[bool]string{true: "▌", false: ""}[m.mcpPage.filtering] + } + rows := m.mcpRows() + m.clampMCPCursor(rows) + height := m.mcpListHeight() + page := &m.mcpPage + if page.cursor < page.offset { + page.offset = page.cursor + } + if page.cursor >= page.offset+height { + page.offset = page.cursor - height + 1 + } + page.offset = max(0, min(page.offset, max(0, len(rows)-height))) + end := min(len(rows), page.offset+height) + statusWidth := mcpStatusWidth(rows[page.offset:end], width) + for index := page.offset; index < end; index++ { + lines = append(lines, ansi.Truncate(m.renderMCPRow(rows[index], index == page.cursor, statusWidth), width, "…")) + } + for len(lines) < height+3 { + lines = append(lines, "") + } + footer := "↑/↓ move · Enter details · / filter · m marketplace · Esc back to chat" + if page.message != "" { + footer = sanitizeTerminalText(page.message) + } + lines = append(lines, "", faint(m.noColor, ansi.Truncate(footer, width, "…"))) + return strings.Join(lines, "\n") +} + +func mcpStatusWidth(rows []mcpRow, width int) int { + statusWidth := 24 + for _, row := range rows { + if row.kind == mcpRowServer { + statusWidth = max(statusWidth, ansi.StringWidth(mcpStatusLabel(*row.view))) + } + } + return min(statusWidth, max(24, width-2-28-1-16)) +} + +func (m *TUIModel) renderMCPRow(row mcpRow, selected bool, statusWidth int) string { + prefix := " " + if selected { + prefix = "❯ " + } + switch row.kind { + case mcpRowHeader: + return m.mcpStyle(lipgloss.NewStyle().Bold(true), row.text) + case mcpRowNote: + return " " + faint(m.noColor, row.text) + case mcpRowLoadMore: + return prefix + row.text + case mcpRowCloud: + name := sanitizeTerminalText(row.cloud.Name) + return prefix + padRight(name, 28) + " " + padRight(cloudStatusLabel(*row.cloud), statusWidth) + " " + faint(m.noColor, sanitizeTerminalText(row.cloud.ProviderType)) + case mcpRowServer: + view := row.view + name := sanitizeTerminalText(view.Server.Name) + origin := sanitizeTerminalText(view.Server.Config.Masked().Origin()) + if view.Server.Scope == mcp.ScopeBuiltin { + origin = "local tools" + } + if view.Server.Scope == mcp.ScopeImported { + origin += " · " + view.Server.SourceLabel + } + if view.Server.Scope == mcp.ScopeLocal { + origin += " · local" + } + return prefix + padRight(name, 28) + " " + padRight(mcpStatusLabel(*view), statusWidth) + " " + faint(m.noColor, origin) + } + return "" +} + +func padRight(text string, width int) string { + text = ansi.Truncate(text, width, "…") + if gap := width - ansi.StringWidth(text); gap > 0 { + return text + strings.Repeat(" ", gap) + } + return text +} + +func (m *TUIModel) mcpDetailView() string { + width := max(20, m.width-1) + bold := lipgloss.NewStyle().Bold(true) + view, cloud := m.detailView() + lines := []string{} + switch { + case cloud != nil: + lines = append(lines, + m.mcpStyle(bold, sanitizeTerminalText(cloud.Name)), + "Provider: "+sanitizeTerminalText(cloud.ProviderType), + "Status: "+cloudStatusLabel(*cloud), + "", + "Cloud connector: permissions come from Tool access on rootly.com", + ) + if cloud.ManageURL != "" { + lines = append(lines, "Manage: "+sanitizeTerminalText(cloud.ManageURL)) + } + lines = append(lines, "", faint(m.noColor, "Esc back")) + return strings.Join(lines, "\n") + case view == nil: + return "This server is no longer configured.\n\n" + faint(m.noColor, "Esc back") + } + if view.Server.Scope == mcp.ScopeBuiltin { + return m.builtinDetailView(view, width, bold) + } + config := view.Server.Config.Masked() + source := displayHomePath(view.Server.Path) + if view.Server.Scope == mcp.ScopeImported { + source = view.Server.SourceLabel + " · " + source + " (read-only)" + } + lines = append(lines, + m.mcpStyle(bold, sanitizeTerminalText(view.Server.Name))+" "+mcpStatusLabel(*view), + faint(m.noColor, sanitizeTerminalText(config.Transport()+" · "+config.Origin())), + faint(m.noColor, "Source: "+sanitizeTerminalText(source)), + ) + if pairs := mcp.MaskedPairs(config.Env, "="); len(pairs) > 0 { + lines = append(lines, faint(m.noColor, "Env: "+sanitizeTerminalText(strings.Join(pairs, ", ")))) + } + if pairs := mcp.MaskedPairs(config.Headers, ": "); len(pairs) > 0 { + lines = append(lines, faint(m.noColor, "Headers: "+sanitizeTerminalText(strings.Join(pairs, ", ")))) + } + if view.Detail != "" { + lines = append(lines, sanitizeTerminalText(view.Detail)) + } + if m.mcpPage.loginURL != "" { + lines = append(lines, "If your browser didn't open, visit: "+sanitizeTerminalText(m.mcpPage.loginURL)) + } + lines = append(lines, "", "[space] on/off for this session [a] approve [l] log in [r] reconnect", "") + tools := m.detailTools(view) + header := fmt.Sprintf("Tools (%d)", len(view.Tools)) + if m.mcpPage.detailFiltering || m.mcpPage.detailFilter != "" { + header += fmt.Sprintf(" · filter: %s%s · %d match", sanitizeTerminalText(m.mcpPage.detailFilter), map[bool]string{true: "▌", false: ""}[m.mcpPage.detailFiltering], len(tools)) + } + lines = append(lines, m.mcpStyle(bold, header)) + height := max(1, m.height-len(lines)-3) + page := &m.mcpPage + page.detailCursor = min(page.detailCursor, max(0, len(tools)-1)) + if page.detailCursor < page.detailOffset { + page.detailOffset = page.detailCursor + } + if page.detailCursor >= page.detailOffset+height { + page.detailOffset = page.detailCursor - height + 1 + } + page.detailOffset = max(0, min(page.detailOffset, max(0, len(tools)-height))) + end := min(len(tools), page.detailOffset+height) + for index := page.detailOffset; index < end; index++ { + tool := tools[index] + prefix := " " + if index == page.detailCursor { + prefix = "❯ " + } + mark := "write" + if tool.ReadOnly { + mark = "read-only" + } + description := strings.Join(strings.Fields(sanitizeTerminalText(mcp.MaskCommand(tool.Description))), " ") + lines = append(lines, ansi.Truncate(prefix+padRight(sanitizeTerminalText(tool.Name), 32)+" "+padRight(mark, 10)+" "+faint(m.noColor, description), width, "…")) + } + if len(view.Tools) == 0 { + lines = append(lines, faint(m.noColor, " No tools yet")) + } + footer := "↑/↓ move · / filter tools · Esc back" + if page.message != "" { + footer = sanitizeTerminalText(page.message) + } + lines = append(lines, "", faint(m.noColor, ansi.Truncate(footer, width, "…"))) + return strings.Join(lines, "\n") +} + +func (m *TUIModel) builtinDetailView(view *mcp.ServerView, width int, bold lipgloss.Style) string { + lines := []string{ + m.mcpStyle(bold, "Built-in tools") + " " + mcpStatusLabel(*view), + faint(m.noColor, "Filesystem access is confined to the current working directory."), + "", + "[space] enable/disable selected tool [bash] asks unless read-only or allowed in ~/.rootly-cli/rules", + } + tools := m.detailTools(view) + header := fmt.Sprintf("Tools (%d)", len(view.Tools)) + if m.mcpPage.detailFiltering || m.mcpPage.detailFilter != "" { + header += fmt.Sprintf(" · filter: %s%s · %d match", sanitizeTerminalText(m.mcpPage.detailFilter), map[bool]string{true: "▌", false: ""}[m.mcpPage.detailFiltering], len(tools)) + } + lines = append(lines, m.mcpStyle(bold, header)) + height := max(1, m.height-len(lines)-3) + page := &m.mcpPage + page.detailCursor = min(page.detailCursor, max(0, len(tools)-1)) + if page.detailCursor < page.detailOffset { + page.detailOffset = page.detailCursor + } + if page.detailCursor >= page.detailOffset+height { + page.detailOffset = page.detailCursor - height + 1 + } + page.detailOffset = max(0, min(page.detailOffset, max(0, len(tools)-height))) + end := min(len(tools), page.detailOffset+height) + for index := page.detailOffset; index < end; index++ { + tool := tools[index] + prefix := " " + if index == page.detailCursor { + prefix = "❯ " + } + mark := "off" + if tool.Enabled { + mark = "enabled" + } + if tool.Name == "local_bash" { + mark += " · approval" + } + description := strings.Join(strings.Fields(sanitizeTerminalText(mcp.MaskCommand(tool.Description))), " ") + lines = append(lines, ansi.Truncate(prefix+padRight(sanitizeTerminalText(tool.Name), 32)+" "+padRight(mark, 16)+" "+faint(m.noColor, description), width, "…")) + } + if len(view.Tools) == 0 { + lines = append(lines, faint(m.noColor, " No built-in tools available")) + } + footer := "↑/↓ move · space toggle · / filter tools · Esc back" + if page.message != "" { + footer = sanitizeTerminalText(page.message) + } + lines = append(lines, "", faint(m.noColor, ansi.Truncate(footer, width, "…"))) + return strings.Join(lines, "\n") +} + +// mcpStatusLine is the pinned line about local servers above the prompt. +func (m *TUIModel) mcpStatusLine() string { + if m.mcp == nil { + return "" + } + if connecting := m.mcp.Connecting(); connecting > 0 { + noun := "servers" + if connecting == 1 { + noun = "server" + } + return faint(m.noColor, fmt.Sprintf("MCP · connecting %d %s…", connecting, noun)) + } + attention := 0 + for _, view := range m.mcpPage.views { + if view.Server.Scope != mcp.ScopeBuiltin && !view.SessionOff && + (view.Status == mcp.StatusError || view.Status == mcp.StatusNeedsAuth || view.Status == mcp.StatusPending) { + attention++ + } + } + if attention == 0 { + return "" + } + noun := "servers need" + if attention == 1 { + noun = "server needs" + } + return faint(m.noColor, fmt.Sprintf("MCP · %d %s attention · /mcp", attention, noun)) +} + +func (m *TUIModel) updateLocalMCP(message tea.Msg) (tea.Cmd, bool) { + switch msg := message.(type) { + case mcpChangedMsg: + m.refreshMCPViews() + m.reportUnreachableMCP() + return m.watchMCP(), true + case mcpStartupCheckMsg: + m.refreshMCPViews() + m.reportUnreachableMCP() + return nil, true + case cloudConnectorsMsg: + m.handleCloudConnectors(msg) + return nil, true + case mcpLoginMsg: + return m.handleMCPLogin(msg), true + case mcpPromptMsg: + return m.handleMCPPrompt(msg), true + case mentionsReadyMsg: + return m.handleMentionsReady(msg), true + case marketLoadedMsg: + m.handleMarketLoaded(msg) + return nil, true + case localToolDoneMsg: + if msg.err != nil && msg.turn == m.turn && m.state == tuiStreaming { + return m.failLocalToolDelivery(msg), true + } + var cmds []tea.Cmd + if msg.failure != "" && msg.turn == m.turn { + cmds = append(cmds, m.println(m.activityLine(msg.failure))) + } + if msg.credHint && !m.credHintShown { + m.credHintShown = true + cmds = append(cmds, m.println(faint(m.noColor, localCredsHint))) + } + if cmd := m.handleLocalToolDone(msg); cmd != nil { + cmds = append(cmds, cmd) + } + return tea.Batch(cmds...), true + } + return nil, false +} + +// failLocalToolDelivery ends the turn when a local result can't reach the +// server, as print mode does: the server would otherwise wait for it forever. +// finishTurn cancels the stream and drops the rest of the local queue. +// The call's own activity keeps its real local outcome first, so a write that +// ran isn't shown as failed and retried. +func (m *TUIModel) failLocalToolDelivery(msg localToolDoneMsg) tea.Cmd { + if msg.failure != "" { + m.println(m.activityLine(msg.failure)) + } + if msg.isError { + m.activities.discardPendingIf(msg.activity) + } else { + for _, line := range m.activities.completePendingIf(msg.activity) { + m.println(m.activityLine(line)) + } + } + _, cmd := m.finishTurn(fmt.Errorf("✗ %s: result not delivered to Rootly, turn stopped: %w", msg.tool, msg.err)) + return cmd +} + +func (m *TUIModel) handleLocalToolDone(msg localToolDoneMsg) tea.Cmd { + if msg.turn != m.turn { + return nil + } + if msg.isError { + m.activities.discardPendingIf(msg.activity) + } + if msg.approval { + m.localApprovalRunning = false + } + if msg.readOnly && m.localReadOnlyRunning > 0 { + m.localReadOnlyRunning-- + } + return m.startNextLocalTool() +} + +func (m *TUIModel) updateOtherMessage(message tea.Msg) (tea.Model, tea.Cmd) { + if m.updateFocus(message) { + return m, nil + } + if cmd, handled := m.updateLocalMCP(message); handled { + return m, cmd + } + return m.updatePickerMessage(message) +} + +func mcpCounts(view mcp.ServerView) string { + parts := []string{pluralTools(len(view.Tools))} + if count := len(view.Resources); count > 0 { + parts = append(parts, pluralCount(count, "resource")) + } + if count := len(view.Prompts); count > 0 { + parts = append(parts, pluralCount(count, "prompt")) + } + return strings.Join(parts, " · ") +} + +func pluralCount(count int, noun string) string { + if count == 1 { + return "1 " + noun + } + return strconv.Itoa(count) + " " + noun + "s" +} diff --git a/internal/agent/mcp_page_test.go b/internal/agent/mcp_page_test.go new file mode 100644 index 0000000..f991a87 --- /dev/null +++ b/internal/agent/mcp_page_test.go @@ -0,0 +1,665 @@ +package agent + +import ( + "context" + "errors" + "fmt" + "iter" + "strings" + "sync" + "testing" + "time" + + "github.com/charmbracelet/x/ansi" + + tea "charm.land/bubbletea/v2" + rootly "github.com/rootlyhq/rootly-go" + + "github.com/rootlyhq/rootly-cli/internal/api" + "github.com/rootlyhq/rootly-cli/internal/mcp" +) + +type fakeMCPManager struct { + fakeLocalTools + views []mcp.ServerView + builtins []mcp.BuiltinToolInfo + toggled []string + approved []string + builtinToggles []string + problems []string + changes chan struct{} + resources map[string]string + readLimits []int + added []mcp.MarketEntry + promptText string +} + +func (f *fakeMCPManager) Views() []mcp.ServerView { return f.views } +func (f *fakeMCPManager) BuiltinTools() []mcp.BuiltinToolInfo { return f.builtins } +func (f *fakeMCPManager) SetBuiltinEnabled(name string, enabled bool) error { + f.builtinToggles = append(f.builtinToggles, fmt.Sprintf("%s:%t", name, enabled)) + for index := range f.builtins { + if f.builtins[index].Name == name { + f.builtins[index].Enabled = enabled + } + } + return nil +} +func (f *fakeMCPManager) Changes() <-chan struct{} { return f.changes } +func (f *fakeMCPManager) Problems() []string { return f.problems } +func (f *fakeMCPManager) Connecting() int { return 0 } +func (f *fakeMCPManager) Enable(string) error { return nil } +func (f *fakeMCPManager) Reconnect(string) error { return nil } +func (f *fakeMCPManager) ToggleSession(id string) error { + f.toggled = append(f.toggled, id) + return nil +} +func (f *fakeMCPManager) Approve(id string) error { f.approved = append(f.approved, id); return nil } +func (f *fakeMCPManager) Login(context.Context, string, mcp.LoginOptions) error { return nil } + +func (f *fakeMCPManager) AddFromMarket(entry mcp.MarketEntry) (mcp.Server, error) { + for _, view := range f.views { + if view.Server.Name == entry.Name { + return mcp.Server{}, errors.New("an MCP server named " + entry.Name + " already exists") + } + } + server := mcp.Server{Name: entry.Name, Scope: mcp.ScopeUser, Config: entry.Config} + f.added = append(f.added, entry) + f.views = append(f.views, mcp.ServerView{Server: server, Status: mcp.StatusOff}) + return server, nil +} + +func (f *fakeMCPManager) GetPrompt(_ context.Context, serverID, name string, arguments map[string]string) (string, error) { + if !strings.HasSuffix(serverID, ":docs") { + return "", errors.New("not connected") + } + if f.promptText != "" { + return f.promptText, nil + } + return fmt.Sprintf("prompt %s %v", name, arguments), nil +} + +func (f *fakeMCPManager) ReadResource(_ context.Context, serverID, uri string, limit int) (text string, truncated bool, err error) { + f.mu.Lock() + f.readLimits = append(f.readLimits, limit) + f.mu.Unlock() + text, ok := f.resources[serverID[strings.LastIndex(serverID, ":")+1:]+":"+uri] + if !ok { + return "", false, errors.New("no text content") + } + if len(text) > limit { + return text[:limit], true, nil + } + return text, false, nil +} + +type fakeConnectorAPI struct { + fakeTUIAPI + pages []int +} + +func (f *fakeConnectorAPI) ListAIConnectors(_ context.Context, page, _ int) (*api.AIConnectorPage, error) { + f.pages = append(f.pages, page) + count := 12 + if page == 1 { + return &api.AIConnectorPage{Connectors: []api.AIConnector{{ID: "c1", Name: "Datadog", ProviderType: "datadog", AvailableToAgent: true, ToolCount: &count}}, NextPage: 2}, nil + } + return &api.AIConnectorPage{Connectors: []api.AIConnector{{ID: "c2", Name: "Sentry", ProviderType: "sentry"}}}, nil +} + +func newMCPTestModel(t *testing.T, manager *fakeMCPManager, connectors *fakeConnectorAPI) *TUIModel { + t.Helper() + model := NewTUIModel(context.Background(), &fakeStreamer{}, TUIOptions{NoColor: true, API: connectors, MCP: manager}) + model.Update(tea.WindowSizeMsg{Width: 120, Height: 40}) + return model +} + +func runCmd(model *TUIModel, cmd tea.Cmd) { + if cmd == nil { + return + } + if msg := cmd(); msg != nil { + model.Update(msg) + } +} + +func testViews() []mcp.ServerView { + tools := make([]mcp.ToolInfo, 1000) + for index := range tools { + tools[index] = mcp.ToolInfo{Name: "tool_" + strings.Repeat("x", index%5) + string(rune('a'+index%26)), ReadOnly: index%2 == 0} + } + tools[999] = mcp.ToolInfo{Name: "write_file", Description: "Write a file"} + return []mcp.ServerView{ + {Server: mcp.Server{Name: "maestro", Scope: mcp.ScopeProject, Path: "/repo/.mcp.json", Config: mcp.ServerConfig{Command: "maestro", Args: []string{"mcp"}}}, Status: mcp.StatusPending}, + {Server: mcp.Server{Name: "fs", Scope: mcp.ScopeUser, Config: mcp.ServerConfig{Command: "npx", Env: map[string]string{"TOKEN": "secret-value"}}}, Status: mcp.StatusConnected, Tools: tools}, + {Server: mcp.Server{Name: "braintrust", Scope: mcp.ScopeUser, Config: mcp.ServerConfig{Type: "http", URL: "https://api.braintrust.dev/mcp", Headers: map[string]string{"Authorization": "Bearer hidden"}}}, Status: mcp.StatusNeedsAuth}, + {Server: mcp.Server{Name: "sentry", Scope: mcp.ScopeImported, SourceLabel: "Claude Code", Config: mcp.ServerConfig{Type: "http", URL: "https://mcp.sentry.dev/mcp"}}, Status: mcp.StatusOff}, + } +} + +func TestMCPPageOpensGroupsAndEscRestoresChat(t *testing.T) { + manager := &fakeMCPManager{views: testViews(), changes: make(chan struct{})} + connectors := &fakeConnectorAPI{} + model := newMCPTestModel(t, manager, connectors) + model.composer.SetValue("draft question") + before := model.View().Content + + for _, command := range []string{"/connectors", "/mcp"} { + model.composer.SetValue(command) + _, cmd := model.Update(tea.KeyPressMsg{Code: tea.KeyEnter}) + runCmd(model, cmd) + if !model.mcpPage.open { + t.Fatalf("%s did not open the MCP page", command) + } + view := model.View().Content + for _, want := range []string{"Manage MCP servers", "Rootly cloud", "Datadog", "✔ 12 tools", "Load more", "Project MCPs", "‖ pending approval", "User MCPs (~/.rootly-cli/mcp.json)", "✔ 1000 tools", "▲ needs authentication", "Imported (read-only)", "○ off", "Esc back to chat"} { + if !strings.Contains(view, want) { + t.Errorf("%s view missing %q:\n%s", command, want, view) + } + } + if strings.Contains(view, "secret-value") || strings.Contains(view, "hidden") { + t.Fatal("MCP page leaks a secret") + } + model.Update(tea.KeyPressMsg{Code: tea.KeyEscape}) + if model.mcpPage.open { + t.Fatal("Esc should close the MCP page") + } + } + model.composer.SetValue("draft question") + if after := model.View().Content; after != before { + t.Fatalf("chat changed after Esc:\nbefore:\n%s\nafter:\n%s", before, after) + } + if len(connectors.pages) != 2 || connectors.pages[1] != 1 { + t.Fatalf("each /mcp visit should refetch page 1, pages = %v", connectors.pages) + } +} + +func TestMCPPageLoadMoreDetailFilterAndActions(t *testing.T) { + manager := &fakeMCPManager{views: testViews(), changes: make(chan struct{})} + connectors := &fakeConnectorAPI{} + model := newMCPTestModel(t, manager, connectors) + runCmd(model, model.openMCPPage()) + + press := func(keys ...tea.KeyPressMsg) { + for _, key := range keys { + _, cmd := model.Update(key) + runCmd(model, cmd) + } + } + down := tea.KeyPressMsg{Code: tea.KeyDown} + enter := tea.KeyPressMsg{Code: tea.KeyEnter} + press(down, enter) + if len(connectors.pages) != 2 || !strings.Contains(model.View().Content, "Sentry") { + t.Fatalf("Load more did not fetch page 2: %v", connectors.pages) + } + + model.mcpPage.cursor = 0 + press(enter) + if view := model.View().Content; !strings.Contains(view, "permissions come from Tool access on rootly.com") { + t.Fatalf("cloud detail:\n%s", view) + } + press(tea.KeyPressMsg{Code: tea.KeyEscape}) + + press(down, down) + row, _ := model.selectedMCPRow() + if row.view == nil || row.view.Server.Name != "maestro" { + t.Fatalf("expected maestro row, got %+v", row) + } + press(tea.KeyPressMsg{Code: 'a', Text: "a"}) + if len(manager.approved) != 1 { + t.Fatal("a should approve") + } + press(down, enter) + view := model.View().Content + if !strings.Contains(view, "Tools (1000)") || !strings.Contains(view, "[space] on/off for this session") || strings.Contains(view, "secret-value") { + t.Fatalf("detail view:\n%s", view) + } + if lines := strings.Count(view, "\n"); lines > 40 { + t.Fatalf("detail list not virtualized: %d lines", lines) + } + press(tea.KeyPressMsg{Code: '/', Text: "/"}) + for _, r := range "write" { + press(tea.KeyPressMsg{Code: r, Text: string(r)}) + } + press(enter) + if view := model.View().Content; !strings.Contains(view, "write_file") || !strings.Contains(view, "1 match") { + t.Fatalf("filtered detail:\n%s", view) + } + press(tea.KeyPressMsg{Code: tea.KeySpace, Text: " "}) + if len(manager.toggled) != 1 || manager.toggled[0] != testViews()[1].Server.ID() { + t.Fatalf("space should toggle fs, got %v", manager.toggled) + } +} + +func TestTUIClientToolWriteNeedsApproval(t *testing.T) { + manager := &fakeMCPManager{changes: make(chan struct{})} + manager.tools = map[string]mcp.ToolRef{"local_fs_write_file": {ServerID: "user:fs", ServerName: "fs", Tool: "write_file"}} + model := newMCPTestModel(t, manager, &fakeConnectorAPI{}) + model.state = tuiStreaming + model.turn = 1 + _, cmd := model.Update(streamMessage{turn: 1, event: Event{Type: EventClientToolCall, ClientToolCall: &ClientToolCall{CallID: "c1", Name: "local_fs_write_file", Arguments: map[string]any{"path": "/tmp/x"}}}}) + _ = cmd + if model.localConfirm == nil || !strings.Contains(model.View().Content, "write_file") { + t.Fatalf("expected a confirm card:\n%s", model.View().Content) + } + _, cmd = model.Update(tea.KeyPressMsg{Code: 'n', Text: "n"}) + if model.localConfirm != nil || cmd == nil { + t.Fatal("n should decline and post a result") + } + if msg, ok := cmd().(localToolDoneMsg); !ok || msg.err != nil { + t.Fatalf("decline result = %#v", msg) + } + if len(manager.calls) != 0 { + t.Fatal("declined tool must not run") + } +} + +func TestTUIFailedLocalToolPrintsCross(t *testing.T) { + manager := &fakeMCPManager{changes: make(chan struct{})} + manager.tools = map[string]mcp.ToolRef{"local_fs_read_text_file": {ServerID: "user:fs", ServerName: "fs", Tool: "read_text_file", ReadOnly: true}} + call := &ClientToolCall{CallID: "c1", Name: "local_fs_read_text_file"} + line := localToolFailure(manager, call, "ENOENT: no such file or directory, open '/tmp/x'\nstack") + if line != "✗ local · fs · read_text_file failed: ENOENT: no such file or directory, open '/tmp/x'" { + t.Fatalf("failure line = %q", line) + } + model := newMCPTestModel(t, manager, &fakeConnectorAPI{}) + model.turn = 1 + model.activities.Handle(activityStatusEvent("Local read file...")) + activity := model.activities.pending + model.Update(localToolDoneMsg{turn: 1, isError: true, activity: activity, failure: line}) + model.handleEvent(activityDoneEvent("complete")) + if !strings.Contains(transcriptText(model), "read_text_file failed: ENOENT") { + t.Fatalf("transcript = %v", model.transcript) + } + if got := strings.Count(transcriptText(model), "✗"); got != 1 { + t.Errorf("failure markers = %d, want exactly one:\n%s", got, transcriptText(model)) + } + if got := strings.Count(transcriptText(model), "✓"); got != 0 { + t.Errorf("success markers after failure = %d, want none:\n%s", got, transcriptText(model)) + } + model.Update(localToolDoneMsg{turn: 0, failure: "✗ stale"}) + if strings.Contains(transcriptText(model), "stale") { + t.Fatal("results from an earlier turn must not print") + } + + successModel := newMCPTestModel(t, manager, &fakeConnectorAPI{}) + successModel.turn = 1 + successModel.activities.Handle(activityStatusEvent("Local read file...")) + successModel.Update(localToolDoneMsg{turn: 1}) + successModel.handleEvent(activityDoneEvent("complete")) + if got := strings.Count(transcriptText(successModel), "✓"); got != 1 { + t.Errorf("success markers = %d, want one:\n%s", got, transcriptText(successModel)) + } +} + +func TestTUIFailedLocalToolPreservesNewerActivity(t *testing.T) { + manager := &fakeMCPManager{ + fakeLocalTools: fakeLocalTools{ + tools: map[string]mcp.ToolRef{ + "local_fs_read_text_file": {ServerID: "user:fs", ServerName: "fs", Tool: "read_text_file", ReadOnly: true}, + }, + output: "path is denied", + }, + changes: make(chan struct{}), + } + model := newMCPTestModel(t, manager, &fakeConnectorAPI{}) + model.turn = 1 + model.handleEvent(activityStatusEvent("Local read file...")) + originalActivity := model.activities.pending + cmd := model.handleClientToolCall(&ClientToolCall{CallID: "c1", Name: "local_fs_read_text_file"}) + if cmd == nil { + t.Fatal("local tool command = nil") + } + message := cmd() + done, ok := message.(localToolDoneMsg) + if !ok { + t.Fatalf("local tool result message = %T, want localToolDoneMsg", message) + } + if done.activity != originalActivity { + t.Fatal("local tool completion did not retain the activity captured when dispatched") + } + + model.handleEvent(activityStatusEvent("Search incidents...")) + newerActivity := model.activities.pending + if newerActivity == nil || newerActivity == originalActivity { + t.Fatal("new activity was not made pending before the tool completion") + } + model.Update(done) + if model.activities.pending != newerActivity { + t.Fatal("failed local tool discarded the newer pending activity") + } + model.handleEvent(activityDoneEvent("complete")) + if !strings.Contains(transcriptText(model), "✓ Search incidents") { + t.Fatalf("newer activity completion was dropped:\n%s", transcriptText(model)) + } +} + +func transcriptText(model *TUIModel) string { + lines := make([]string, 0, len(model.transcript)) + for _, block := range model.transcript { + lines = append(lines, block.text) + } + return strings.Join(lines, "\n") +} + +func TestBuiltinMCPPageTogglesOnlyWithSpace(t *testing.T) { + manager := &fakeMCPManager{ + builtins: []mcp.BuiltinToolInfo{ + {Name: "local_read_file", Description: "Read a file", ReadOnly: true, Enabled: true}, + {Name: "local_bash", Description: "Run bash", Enabled: false}, + }, + changes: make(chan struct{}), + } + model := newMCPTestModel(t, manager, &fakeConnectorAPI{}) + model.refreshMCPViews() + view := &model.mcpPage.views[0] + model.mcpAction(view, "a") + if len(manager.builtinToggles) != 0 { + t.Fatal("built-in tools should not expose the MCP approve action") + } + rows := model.mcpRows() + builtinRow := -1 + for index, row := range rows { + if row.kind == mcpRowServer && row.view.Server.Scope == mcp.ScopeBuiltin { + builtinRow = index + break + } + } + if builtinRow < 0 { + t.Fatal("built-in group row was not present in the MCP list") + } + model.mcpPage.open = true + model.mcpPage.cursor = builtinRow + before := append([]mcp.BuiltinToolInfo(nil), manager.builtins...) + model.handleMCPPageKey(tea.KeyPressMsg{Code: tea.KeySpace, Text: " "}) + if len(manager.builtinToggles) != 0 { + t.Fatalf("space on the built-in list row toggled tools: %v", manager.builtinToggles) + } + for index, tool := range manager.builtins { + if tool.Enabled != before[index].Enabled { + t.Fatalf("space on the built-in list row changed %s", tool.Name) + } + } + model.handleMCPPageKey(tea.KeyPressMsg{Code: tea.KeyEnter}) + model.handleMCPPageKey(tea.KeyPressMsg{Code: tea.KeySpace, Text: " "}) + if got := strings.Join(manager.builtinToggles, ","); got != "local_read_file:false" { + t.Fatalf("space toggle calls = %s", got) + } + if got := mcpStatusLabel(model.mcpPage.views[0]); got != "0/2 enabled" { + t.Fatalf("built-in summary status = %q", got) + } +} + +func TestTUILocalToolDispatchSerializesWritesWithReadOnlyCalls(t *testing.T) { + manager := &fakeMCPManager{changes: make(chan struct{})} + manager.tools = map[string]mcp.ToolRef{ + "local_read_one": {ServerID: "builtin", ServerName: "Built-in", Tool: "read_one", ReadOnly: true}, + "local_write_one": {ServerID: "builtin", ServerName: "Built-in", Tool: "write_one"}, + "local_read_two": {ServerID: "builtin", ServerName: "Built-in", Tool: "read_two", ReadOnly: true}, + } + model := newMCPTestModel(t, manager, &fakeConnectorAPI{}) + model.turn = 1 + + readCommand := model.handleClientToolCall(&ClientToolCall{CallID: "read-1", Name: "local_read_one"}) + if readCommand == nil || model.localReadOnlyRunning != 1 { + t.Fatal("read-only tool should start concurrently") + } + readDone, ok := readCommand().(localToolDoneMsg) + if !ok || !readDone.readOnly { + t.Fatalf("read-only completion = %#v", readDone) + } + _ = model.handleClientToolCall(&ClientToolCall{CallID: "write-1", Name: "local_write_one"}) + _ = model.handleClientToolCall(&ClientToolCall{CallID: "read-2", Name: "local_read_two"}) + if model.localConfirm != nil || len(model.localToolQueue) != 2 { + t.Fatal("write and subsequent read should wait behind the active read") + } + if cmd := model.handleLocalToolDone(readDone); cmd != nil { + t.Fatal("write confirmation should be presented before starting queued calls") + } + if model.localConfirm == nil || model.localConfirm.call.Name != "local_write_one" { + t.Fatal("queued write was not presented for approval after the read completed") + } + writeCommand := model.resolveLocalConfirmation(true) + writeDone, ok := writeCommand().(localToolDoneMsg) + if !ok || !writeDone.approval { + t.Fatalf("approved write completion = %#v", writeDone) + } + next := model.handleLocalToolDone(writeDone) + if next == nil { + t.Fatal("queued read-only call did not start after the approved write completed") + } + runCmd(model, next) + if got := strings.Join(manager.calls, ","); got != "local_read_one,local_write_one,local_read_two" { + t.Fatalf("local tool order = %s", got) + } + if model.localReadOnlyRunning != 0 || model.localApprovalRunning { + t.Fatal("local tool scheduler did not return to idle") + } +} + +type fakeApprovalMemory struct { + fakeMCPManager + options mcp.ApprovalOptions + remembered []string + allowed map[string]bool +} + +func (f *fakeApprovalMemory) NeedsApproval(name string, args map[string]any) bool { + return !f.allowed[name] && f.fakeMCPManager.NeedsApproval(name, args) +} + +func (f *fakeApprovalMemory) ApprovalOptions(string, map[string]any) mcp.ApprovalOptions { + return f.options +} + +func (f *fakeApprovalMemory) RememberApproval(name string, _ map[string]any, persist bool) error { + f.remembered = append(f.remembered, fmt.Sprintf("%s:%t", name, persist)) + f.allowed[name] = true + return nil +} + +func TestTUILocalApprovalKeysRememberOnlyWhenOffered(t *testing.T) { + for _, test := range []struct { + name string + options mcp.ApprovalOptions + key string + remembered string + footer string + transcript string + }{ + {name: "once", options: mcp.ApprovalOptions{Session: true, Persist: true, Prefix: []string{"kubectl", "rollout"}}, key: "y", + footer: "a allow `kubectl rollout` this session · p always allow `kubectl rollout`"}, + {name: "session prefix", options: mcp.ApprovalOptions{Session: true, Persist: true, Prefix: []string{"kubectl", "rollout"}}, key: "a", + remembered: "local_kubectl:false", transcript: "Allowing `kubectl rollout` for this session"}, + {name: "permanent prefix", options: mcp.ApprovalOptions{Session: true, Persist: true, Prefix: []string{"kubectl", "rollout"}}, key: "p", + remembered: "local_kubectl:true", transcript: "Always allowing `kubectl rollout` (saved to ~/.rootly-cli/rules/default.rules)"}, + {name: "mcp tool session", options: mcp.ApprovalOptions{Session: true}, key: "a", + remembered: "local_kubectl:false", footer: "y once · a allow this tool this session · n/Esc cancel", transcript: "Allowing this tool for this session"}, + {name: "mcp tool cannot persist", options: mcp.ApprovalOptions{Session: true}, key: "p"}, + {name: "complex bash is once only", key: "a", footer: "y approve · n/Esc cancel"}, + {name: "complex bash cannot persist", key: "p"}, + } { + t.Run(test.name, func(t *testing.T) { + manager := &fakeApprovalMemory{options: test.options, allowed: map[string]bool{}} + manager.changes = make(chan struct{}) + manager.tools = map[string]mcp.ToolRef{"local_kubectl": {ServerID: "builtin", ServerName: "Built-in", Tool: "kubectl"}} + model := NewTUIModel(context.Background(), &fakeStreamer{}, TUIOptions{NoColor: true, API: &fakeConnectorAPI{}, MCP: manager}) + model.Update(tea.WindowSizeMsg{Width: 200, Height: 40}) + model.turn = 1 + model.state = tuiStreaming + call := &ClientToolCall{CallID: "c1", Name: "local_kubectl", Arguments: map[string]any{"args": []any{"rollout", "undo"}}} + if cmd := model.handleClientToolCall(call); cmd != nil || model.localConfirm == nil { + t.Fatal("approval card was not shown") + } + if view := strings.Join(strings.Fields(ansi.Strip(model.View().Content)), " "); test.footer != "" && !strings.Contains(view, test.footer) { + t.Errorf("card footer missing %q in %q", test.footer, view) + } + _, cmd := model.Update(tea.KeyPressMsg{Code: rune(test.key[0]), Text: test.key}) + ran := test.key == "y" || test.remembered != "" + if (cmd != nil) != ran || (model.localConfirm == nil) != ran { + t.Fatalf("key %q: ran=%v cmd=%v confirm=%v", test.key, ran, cmd != nil, model.localConfirm != nil) + } + if got := strings.Join(manager.remembered, ","); got != test.remembered { + t.Errorf("remembered = %q, want %q", got, test.remembered) + } + if test.transcript != "" && !strings.Contains(ansi.Strip(model.Transcript()), test.transcript) { + t.Errorf("transcript missing %q: %q", test.transcript, ansi.Strip(model.Transcript())) + } + if !ran { + return + } + runCmd(model, cmd) + model.turn = 1 + second := model.handleClientToolCall(&ClientToolCall{CallID: "c2", Name: "local_kubectl"}) + if remembered := test.remembered != ""; (second != nil) != remembered { + t.Errorf("second call auto-ran=%v, want %v", second != nil, remembered) + } + }) + } +} + +func TestTUIRememberedWritesStillRunOneAtATime(t *testing.T) { + manager := &fakeApprovalMemory{allowed: map[string]bool{"local_apply": true}} + manager.changes = make(chan struct{}) + manager.tools = map[string]mcp.ToolRef{ + "local_apply": {ServerID: "builtin", ServerName: "Built-in", Tool: "apply"}, + "local_read": {ServerID: "builtin", ServerName: "Built-in", Tool: "read", ReadOnly: true}, + } + model := NewTUIModel(context.Background(), &fakeStreamer{}, TUIOptions{NoColor: true, API: &fakeConnectorAPI{}, MCP: manager}) + model.Update(tea.WindowSizeMsg{Width: 200, Height: 40}) + model.turn = 1 + + first := model.handleClientToolCall(&ClientToolCall{CallID: "apply-1", Name: "local_apply"}) + if first == nil || !model.localApprovalRunning || model.localReadOnlyRunning != 0 || model.localConfirm != nil { + t.Fatal("a remembered write should start alone without a card") + } + if cmd := model.handleClientToolCall(&ClientToolCall{CallID: "apply-2", Name: "local_apply"}); cmd != nil { + t.Fatal("a second remembered write started while the first ran") + } + if cmd := model.handleClientToolCall(&ClientToolCall{CallID: "read-1", Name: "local_read"}); cmd != nil { + t.Fatal("a read started while a write ran") + } + firstDone := first().(localToolDoneMsg) + second := model.handleLocalToolDone(firstDone) + if second == nil || !model.localApprovalRunning || model.localReadOnlyRunning != 0 || len(model.localToolQueue) != 1 { + t.Fatal("the queued write should run next, alone") + } + secondDone, ok := second().(localToolDoneMsg) + if !ok || !secondDone.approval { + t.Fatalf("second write was batched: %#v", secondDone) + } + third := model.handleLocalToolDone(secondDone) + if third == nil || model.localApprovalRunning || model.localReadOnlyRunning != 1 { + t.Fatal("the read should start after the writes") + } + runCmd(model, third) + if got := strings.Join(manager.calls, ","); got != "local_apply,local_apply,local_read" { + t.Fatalf("local tool order = %s", got) + } +} + +func TestTUIFailedToolResultDeliveryEndsTheTurn(t *testing.T) { + for _, test := range []struct { + name, toolOutput, wantActivity string + }{ + {name: "tool ran", wantActivity: "✓ Local read one"}, + {name: "tool failed", toolOutput: "boom", wantActivity: "✗ local · Built-in · read_one failed: boom"}, + } { + t.Run(test.name, func(t *testing.T) { + manager := &fakeMCPManager{changes: make(chan struct{})} + manager.output = test.toolOutput + manager.tools = map[string]mcp.ToolRef{ + "local_read_one": {ServerID: "builtin", ServerName: "Built-in", Tool: "read_one", ReadOnly: true}, + "local_write_one": {ServerID: "builtin", ServerName: "Built-in", Tool: "write_one"}, + } + streamer := &twoTurnStreamer{firstStopped: make(chan struct{})} + poster := &switchableToolResultPoster{err: errors.New("post failed \x1b]52;c;bad\x07")} + model := NewTUIModel(context.Background(), streamer, TUIOptions{NoColor: true, API: &fakeConnectorAPI{}, MCP: manager, ToolResults: poster}) + model.Update(tea.WindowSizeMsg{Width: 200, Height: 40}) + model.submitPrompt("first") + streamCtx := model.streamCtx + + model.activities.Handle(activityStatusEvent("Local read one...")) + read := model.handleClientToolCall(&ClientToolCall{CallID: "read-1", Name: "local_read_one"}) + model.handleClientToolCall(&ClientToolCall{CallID: "write-1", Name: "local_write_one"}) + if read == nil || len(model.localToolQueue) != 1 { + t.Fatal("expected the read to run and the write to queue") + } + runCmd(model, read) + if model.state != tuiIdle || model.stream != nil || streamCtx.Err() == nil { + t.Fatal("a failed result delivery should cancel the stream and end the turn") + } + select { + case <-streamer.firstStopped: + case <-time.After(2 * time.Second): + t.Fatal("the first stream kept running after the failed delivery") + } + if model.localConfirm != nil || len(model.localToolQueue) != 0 || model.localApprovalRunning || model.localReadOnlyRunning != 0 { + t.Fatal("local tool state was not reset") + } + if got := strings.Join(manager.calls, ","); got != "local_read_one" { + t.Fatalf("queued tool ran after the failed delivery: %s", got) + } + transcript := transcriptText(model) + if strings.Count(transcript, "result not delivered") != 1 || !strings.Contains(transcript, "✗ local · Built-in · read_one: result not delivered") || strings.Contains(transcript, "\x1b") { + t.Fatalf("want one sanitized error line naming the tool:\n%s", transcript) + } + if !strings.Contains(transcript, test.wantActivity) { + t.Fatalf("the tool's own outcome %q is missing:\n%s", test.wantActivity, transcript) + } + wantCrosses := 1 + if test.toolOutput != "" { + wantCrosses = 2 + } + if got := strings.Count(transcript, "✗"); got != wantCrosses { + t.Fatalf("red ✗ lines = %d, want %d:\n%s", got, wantCrosses, transcript) + } + + poster.err = nil + model.submitPrompt("second") + for range 10 { + if model.state != tuiStreaming { + break + } + runCmd(model, model.waitForStream()) + } + if model.state != tuiIdle || model.lastAnswer != "second answer" || model.status != "" { + t.Fatalf("the next prompt did not stream normally:\n%s", transcriptText(model)) + } + }) + } +} + +type switchableToolResultPoster struct { + err error +} + +func (p *switchableToolResultPoster) PostAiChatToolResult(context.Context, api.ToolResult) error { + return p.err +} + +// twoTurnStreamer holds its first turn open, as Rails does while it waits for +// a local tool result, and answers the second. +type twoTurnStreamer struct { + mu sync.Mutex + calls int + firstStopped chan struct{} +} + +func (s *twoTurnStreamer) StreamAiChatEvents(ctx context.Context, _ *rootly.StreamAiChatParams) (iter.Seq2[rootly.AiChatStreamEvent, error], error) { + s.mu.Lock() + s.calls++ + first := s.calls == 1 + s.mu.Unlock() + if first { + return func(func(rootly.AiChatStreamEvent, error) bool) { + <-ctx.Done() + close(s.firstStopped) + }, nil + } + return func(yield func(rootly.AiChatStreamEvent, error) bool) { + if yield(eventOf(rootly.AiChatStreamEventText, `{"content":"second answer"}`), nil) { + yield(eventOf(rootly.AiChatStreamEventDone, `{"status":"complete"}`), nil) + } + }, nil +} diff --git a/internal/agent/metric_chart.go b/internal/agent/metric_chart.go new file mode 100644 index 0000000..3eeaa56 --- /dev/null +++ b/internal/agent/metric_chart.go @@ -0,0 +1,410 @@ +package agent + +import ( + "encoding/json" + "fmt" + "math" + "strconv" + "strings" + + "charm.land/lipgloss/v2" + "github.com/charmbracelet/x/ansi" + "github.com/guptarohit/asciigraph" +) + +const metricChartColor = "#7C5CE6" + +type markdownOpts struct { + width int + indent int + noColor bool + streaming bool + webBase string + sessionID string +} + +type markdownChart struct { + token string + indent string + lines []string +} + +type metricChartSpec struct { + Type string `json:"type"` + Labels []string `json:"labels"` + Values []*float64 `json:"values"` + Displays []string `json:"displays"` +} + +func prepareAssistantMarkdown(markdown string, opts markdownOpts) (string, []markdownChart) { + markdown = sanitizeTerminalText(markdown) + if opts.width <= 0 { + opts.width = 80 + } + markdown = rewriteDashboardLinks(markdown, opts.webBase, opts.sessionID) + + lines := strings.Split(markdown, "\n") + prepared := make([]string, 0, len(lines)) + charts := make([]markdownChart, 0) + usedTokens := make(map[string]struct{}) + var closingFence string + + for index := 0; index < len(lines); index++ { + line := lines[index] + if closingFence != "" { + prepared = append(prepared, line) + if isClosingMarkdownFence(line, closingFence) { + closingFence = "" + } + continue + } + + if isMetricChartFence(line) { + end := findMetricChartFenceEnd(lines, index+1) + if end < 0 { + if opts.streaming { + token := nextMetricChartToken(markdown, usedTokens, len(charts)) + charts = append(charts, markdownChart{ + token: token, + indent: markdownLineIndent(line), + lines: []string{drawingChartLine(opts.noColor)}, + }) + prepared = append(prepared, markdownLineIndent(line)+token) + break + } + prepared = append(prepared, lines[index:]...) + break + } + + chart, valid := renderMetricChart(strings.Join(lines[index+1:end], "\n"), opts) + if !valid { + prepared = append(prepared, lines[index:end+1]...) + } else { + token := nextMetricChartToken(markdown, usedTokens, len(charts)) + indent := markdownLineIndent(line) + charts = append(charts, markdownChart{token: token, indent: indent, lines: chart}) + prepared = append(prepared, indent+token) + } + index = end + continue + } + + prepared = append(prepared, line) + if fence, ok := openingMarkdownFence(line); ok { + closingFence = fence + } + } + + return bulletizeWideTables(strings.Join(prepared, "\n"), opts.width), charts +} + +func expandAssistantMarkdownCharts(markdown string, charts []markdownChart) string { + if len(charts) == 0 { + return markdown + } + + byToken := make(map[string]markdownChart, len(charts)) + for _, chart := range charts { + byToken[chart.token] = chart + } + lines := strings.Split(markdown, "\n") + expanded := make([]string, 0, len(lines)) + for _, line := range lines { + plain := strings.TrimSpace(ansi.Strip(line)) + if chart, ok := byToken[plain]; ok { + indent := chart.indent + stripped := ansi.Strip(line) + if rendered := stripped[:len(stripped)-len(strings.TrimLeft(stripped, " "))]; len(rendered) > len(indent) { + indent = rendered + } + for _, chartLine := range chart.lines { + expanded = append(expanded, indent+chartLine) + } + continue + } + expanded = append(expanded, line) + } + return strings.Join(expanded, "\n") +} + +func renderMetricChart(source string, opts markdownOpts) ([]string, bool) { + var spec metricChartSpec + if err := json.Unmarshal([]byte(source), &spec); err != nil || + len(spec.Values) == 0 || len(spec.Labels) != len(spec.Values) { + return nil, false + } + for _, value := range spec.Values { + if value != nil && (math.IsNaN(*value) || math.IsInf(*value, 0)) { + return nil, false + } + } + + switch spec.Type { + case "pie": + return renderMetricBars(spec, opts, true), true + case "bar": + return renderMetricBars(spec, opts, false), true + case "line", "step": + if metricValueCount(spec.Values) < 2 { + return renderMetricBars(spec, opts, false), true + } + width := clamp(opts.width-12, 20, 60) + values := metricSeries(spec, spec.Type == "step") + var graph string + if spec.Type == "step" { + graph = asciigraph.Plot(expandStepSeries(values, width), asciigraph.Height(8)) + } else { + series, plotWidth := fitLineSeries(values, width) + graph = asciigraph.Plot(series, asciigraph.Height(8), asciigraph.Width(plotWidth)) + } + if !opts.noColor { + graph = lipgloss.NewStyle().Foreground(lipgloss.Color(metricChartColor)).Render(graph) + } + caption := metricChartCaption(spec.Labels, opts.width-opts.indent, opts.noColor) + return append(strings.Split(graph, "\n"), caption), true + default: + return nil, false + } +} + +func renderMetricBars(spec metricChartSpec, opts markdownOpts, pie bool) []string { + labelWidth := 0 + for _, label := range spec.Labels { + labelWidth = max(labelWidth, ansi.StringWidth(sanitizeTerminalText(label))) + } + labelWidth = min(labelWidth, 20) + + sum, maximum := 0.0, 0.0 + for _, value := range spec.Values { + if value == nil { + continue + } + sum += *value + maximum = max(maximum, *value) + } + + suffixes := make([]string, len(spec.Values)) + maxSuffixWidth := 0 + for index, value := range spec.Values { + if value == nil { + suffixes[index] = "–" + } else { + suffix := metricDisplay(spec.Displays, index, *value) + if pie && sum > 0 { + suffix += fmt.Sprintf(" · %d%%", int(math.Round(*value/sum*100))) + } + suffixes[index] = suffix + } + maxSuffixWidth = max(maxSuffixWidth, ansi.StringWidth(suffixes[index])) + } + + labelColumn := labelWidth + 1 + suffixWidth := maxSuffixWidth + 1 + availableWidth := opts.width - labelColumn - suffixWidth - opts.indent + barWidth := clamp(availableWidth, 10, 40) + + lines := make([]string, len(spec.Values)) + for index, value := range spec.Values { + label := ansi.Truncate(sanitizeTerminalText(spec.Labels[index]), labelWidth, "…") + label = ansi.Truncate(label, labelWidth, "") + label += strings.Repeat(" ", max(0, labelWidth-ansi.StringWidth(label))) + bar := "" + if value != nil && *value > 0 && maximum > 0 { + bar = metricBar(*value/maximum, barWidth) + if !opts.noColor { + bar = lipgloss.NewStyle().Foreground(lipgloss.Color(metricChartColor)).Render(bar) + } + } + line := label + " " + bar + " " + suffixes[index] + lines[index] = strings.TrimRight(line, " ") + } + return lines +} + +func metricSeries(spec metricChartSpec, carryNulls bool) []float64 { + values := make([]float64, 0, len(spec.Values)) + var previous float64 + hasPrevious := false + for _, value := range spec.Values { + if value == nil { + if carryNulls && hasPrevious { + values = append(values, previous) + } else if !carryNulls { + values = append(values, math.NaN()) + } + continue + } + previous = *value + hasPrevious = true + values = append(values, previous) + } + return values +} + +func metricValueCount(values []*float64) int { + count := 0 + for _, value := range values { + if value != nil { + count++ + } + } + return count +} + +// fitLineSeries keeps every measurement on a plotted column: asciigraph's +// interpolation only hits the original points when the plot width is +// (n-1)*k+1. Longer series keep their exact endpoints and bucket the +// interior by max so peaks survive. +func fitLineSeries(values []float64, width int) (series []float64, plotWidth int) { + width = max(3, width) + if len(values) < 2 { + return values, max(1, len(values)) + } + if len(values) <= width { + step := (width - 1) / (len(values) - 1) + return values, (len(values)-1)*step + 1 + } + interior := values[1 : len(values)-1] + columns := width - 2 + series = make([]float64, 0, width) + series = append(series, values[0]) + for index := range columns { + start := index * len(interior) / columns + end := max(start+1, (index+1)*len(interior)/columns) + peak := math.NaN() + for _, value := range interior[start:end] { + if !math.IsNaN(value) && (math.IsNaN(peak) || value > peak) { + peak = value + } + } + series = append(series, peak) + } + series = append(series, values[len(values)-1]) + return series, width +} + +func expandStepSeries(values []float64, width int) []float64 { + if len(values) == 0 { + return nil + } + width = max(1, width) + expanded := make([]float64, width) + for index := range expanded { + position := index * len(values) / width + if len(values) > width && width > 1 { + position = int(math.Round(float64(index) * float64(len(values)-1) / float64(width-1))) + } + expanded[index] = values[position] + } + return expanded +} + +func metricChartCaption(labels []string, width int, noColor bool) string { + if len(labels) == 0 { + return "" + } + caption := sanitizeTerminalText(labels[0]) + " … " + sanitizeTerminalText(labels[len(labels)-1]) + caption = ansi.Truncate(caption, max(1, width), "…") + if noColor { + return caption + } + return lipgloss.NewStyle().Faint(true).Render(caption) +} + +func metricBar(ratio float64, width int) string { + if ratio <= 0 || width <= 0 { + return "" + } + fractions := []string{"▏", "▎", "▍", "▌", "▋", "▊", "▉"} + eighths := int(math.Round(min(1, ratio) * float64(width*8))) + if eighths == 0 { + eighths = 1 + } + full := eighths / 8 + remainder := eighths % 8 + bar := strings.Repeat("█", full) + if remainder > 0 { + bar += fractions[remainder-1] + } + return bar +} + +func metricDisplay(displays []string, index int, value float64) string { + if index >= len(displays) || strings.TrimSpace(displays[index]) == "" { + return strconv.FormatFloat(value, 'f', -1, 64) + } + return sanitizeTerminalText(displays[index]) +} + +func drawingChartLine(noColor bool) string { + line := "◌ drawing chart…" + if noColor { + return line + } + return lipgloss.NewStyle().Faint(true).Render(line) +} + +func findMetricChartFenceEnd(lines []string, start int) int { + for index := start; index < len(lines); index++ { + if strings.TrimSpace(lines[index]) == "```" { + return index + } + } + return -1 +} + +func isMetricChartFence(line string) bool { + if _, ok := openingMarkdownFence(line); !ok { + return false + } + return strings.TrimSpace(strings.TrimLeft(line, " ")) == "```metric-chart" +} + +func openingMarkdownFence(line string) (string, bool) { + trimmed := strings.TrimLeft(line, " ") + if len(line)-len(trimmed) > 3 || len(trimmed) < 3 { + return "", false + } + character := trimmed[0] + if character != '`' && character != '~' { + return "", false + } + count := 0 + for count < len(trimmed) && trimmed[count] == character { + count++ + } + if count < 3 { + return "", false + } + return strings.Repeat(string(character), count), true +} + +func markdownLineIndent(line string) string { + trimmed := strings.TrimLeft(line, " \t") + return line[:len(line)-len(trimmed)] +} + +func isClosingMarkdownFence(line, opening string) bool { + trimmed := strings.TrimSpace(line) + if len(trimmed) < len(opening) || !strings.HasPrefix(trimmed, opening) { + return false + } + return strings.TrimSpace(strings.TrimPrefix(trimmed, opening)) == "" +} + +func nextMetricChartToken(markdown string, used map[string]struct{}, index int) string { + for { + token := fmt.Sprintf("RLYCHART%d", index) + if !strings.Contains(markdown, token) { + if _, exists := used[token]; !exists { + used[token] = struct{}{} + return token + } + } + index++ + } +} + +func clamp(value, minimum, maximum int) int { + return min(max(value, minimum), maximum) +} diff --git a/internal/agent/metric_chart_test.go b/internal/agent/metric_chart_test.go new file mode 100644 index 0000000..347dbfc --- /dev/null +++ b/internal/agent/metric_chart_test.go @@ -0,0 +1,319 @@ +package agent + +import ( + "bytes" + "fmt" + "slices" + "strconv" + "strings" + "testing" + + "github.com/charmbracelet/x/ansi" + "github.com/google/uuid" +) + +const metricPieSource = "```metric-chart\n" + + `{"type":"pie","label":"Incident count","unit":"count","dimension":"Severity","labels":["SEV2","SEV3","SEV1"],"values":[13,11,1],"displays":["13","11","1"],"options":[{"color":"#7C5CE6"}]}` + + "\n```" + +func TestPrepareAssistantMarkdownRendersPieChart(t *testing.T) { + for _, noColor := range []bool{false, true} { + prepared, charts := prepareAssistantMarkdown(metricPieSource, markdownOpts{width: 100, noColor: noColor}) + got := expandAssistantMarkdownCharts(prepared, charts) + for _, text := range []string{"SEV2", "13 · 52%", "SEV3", "11 · 44%", "SEV1", "1 · 4%"} { + if !strings.Contains(ansi.Strip(got), text) { + t.Errorf("pie output is missing %q: %q", text, got) + } + } + if strings.Contains(got, `{"type"`) { + t.Errorf("pie output contains raw JSON: %q", got) + } + for _, line := range strings.Split(got, "\n") { + if width := ansi.StringWidth(line); width > 100 { + t.Errorf("pie line width = %d, want at most 100: %q", width, line) + } + } + if noColor && strings.Contains(got, "\x1b") { + t.Errorf("no-color pie output contains ANSI: %q", got) + } + } +} + +func TestPrepareAssistantMarkdownRendersLineChart(t *testing.T) { + source := "```metric-chart\n" + + `{"type":"line","label":"Incidents","unit":"count","labels":["2026-01-01","2026-01-02","2026-01-03","2026-01-04","2026-01-05"],"values":[2,5,3,7,4],"displays":["2","5","3","7","4"]}` + + "\n```" + prepared, charts := prepareAssistantMarkdown(source, markdownOpts{width: 100, noColor: true}) + got := expandAssistantMarkdownCharts(prepared, charts) + if !strings.Contains(got, "╭") { + t.Errorf("line output does not contain asciigraph axes: %q", got) + } + if !strings.Contains(got, "2026-01-01 … 2026-01-05") { + t.Errorf("line output is missing its date caption: %q", got) + } +} + +func TestPrepareAssistantMarkdownLeavesInvalidMetricFence(t *testing.T) { + source := "```metric-chart\n{not json}\n```" + prepared, charts := prepareAssistantMarkdown(source, markdownOpts{width: 100}) + if prepared != source || len(charts) != 0 { + t.Errorf("invalid metric fence = %q, %d charts; want original code fence", prepared, len(charts)) + } +} + +func TestPrepareAssistantMarkdownLeavesIndentedMetricFenceLiteral(t *testing.T) { + source := " ```metric-chart\n" + + " {\"type\":\"pie\",\"labels\":[\"SEV2\"],\"values\":[13]}\n" + + " ```" + prepared, charts := prepareAssistantMarkdown(source, markdownOpts{width: 100}) + if len(charts) != 0 || !strings.Contains(prepared, `{"type"`) { + t.Errorf("indented fence was converted to a chart: %q, %d charts", prepared, len(charts)) + } +} + +func TestPrepareAssistantMarkdownRendersStepAsVerticalTransitions(t *testing.T) { + source := "```metric-chart\n" + + `{"type":"step","label":"Incidents","unit":"count","labels":["a","b","c","d"],"values":[1,1,5,5]}` + + "\n```" + prepared, charts := prepareAssistantMarkdown(source, markdownOpts{width: 100, noColor: true}) + got := expandAssistantMarkdownCharts(prepared, charts) + cornerColumns := make(map[int]struct{}) + hasHorizontalRun := false + for _, line := range strings.Split(got, "\n") { + if strings.Contains(line, "──") { + hasHorizontalRun = true + } + for column, character := range []rune(line) { + if strings.ContainsRune("╭╮╰╯", character) { + cornerColumns[column] = struct{}{} + } + } + } + if !hasHorizontalRun { + t.Errorf("step output has no horizontal plateau: %q", got) + } + if len(cornerColumns) != 1 { + t.Errorf("step transitions span %d columns, want one vertical jump: %q", len(cornerColumns), got) + } +} + +func TestPrepareAssistantMarkdownUsesNumericDisplayFallback(t *testing.T) { + source := "```metric-chart\n" + + `{"type":"pie","labels":["SEV2","SEV3","SEV1"],"values":[13,11,1]}` + + "\n```" + prepared, charts := prepareAssistantMarkdown(source, markdownOpts{width: 100, noColor: true}) + got := expandAssistantMarkdownCharts(prepared, charts) + for _, want := range []string{"13 · 52%", "11 · 44%", "1 · 4%"} { + if !strings.Contains(got, want) { + t.Errorf("output = %q, want numeric fallback %q", got, want) + } + } +} + +func TestPrepareAssistantMarkdownLineChartLeavesNullGap(t *testing.T) { + source := "```metric-chart\n" + + `{"type":"line","labels":["a","b","c"],"values":[1,null,3]}` + + "\n```" + prepared, charts := prepareAssistantMarkdown(source, markdownOpts{width: 100, noColor: true}) + got := expandAssistantMarkdownCharts(prepared, charts) + if !strings.Contains(got, "╴") || !strings.Contains(got, "╶") { + t.Errorf("line output has no separated null-gap endpoints: %q", got) + } +} + +func TestRewriteDashboardLinks(t *testing.T) { + for _, test := range []struct { + name string + href string + sessionID string + want string + }{ + { + name: "with session", + href: "#AI-add-to-dashboard%3aref", + sessionID: testSessionID, + want: "https://rootly.com/account?chat=" + testSessionID, + }, + { + name: "without session", + href: "#ai-add-to-dashboard:ref", + want: "https://rootly.com/account", + }, + } { + t.Run(test.name, func(t *testing.T) { + source := "[Add to dashboard](<" + test.href + ">)" + prepared, charts := prepareAssistantMarkdown(source, markdownOpts{ + width: 100, noColor: true, webBase: "https://rootly.com", sessionID: test.sessionID, + }) + got := expandAssistantMarkdownCharts(prepared, charts) + if !strings.Contains(got, "[Add to dashboard: open in web]("+test.want+")") { + t.Errorf("rewritten link = %q, want destination %q", got, test.want) + } + if strings.Contains(strings.ToLower(got), "ai-add-to-dashboard") { + t.Errorf("dashboard marker remains in output: %q", got) + } + }) + } +} + +func TestLiveMetricFenceShowsDrawingPlaceholder(t *testing.T) { + model := NewTUIModel(t.Context(), &fakeStreamer{}, TUIOptions{NoColor: true}) + model.state = tuiStreaming + model.answer.WriteString("```metric-chart\n{\"type\":\"pie\",\"values\":[") + + got := ansi.Strip(model.renderLiveAssistantAnswerPart(model.answer.String(), true)) + if !strings.Contains(got, "◌ drawing chart…") { + t.Errorf("live output = %q, want drawing placeholder", got) + } + if strings.Contains(got, `"type"`) || strings.Contains(got, `"values"`) { + t.Errorf("live output exposed partial metric JSON: %q", got) + } +} + +func TestMetricChartsReflowOnTranscriptResize(t *testing.T) { + for _, width := range []int{60, 120} { + t.Run(strconv.Itoa(width), func(t *testing.T) { + model := NewTUIModel(t.Context(), &fakeStreamer{}, TUIOptions{NoColor: true}) + model.width = width + model.appendAssistantTranscriptBlock(metricPieSource, true, false) + model.rebuildTranscriptCache() + if !strings.Contains(model.transcriptCache, "13 · 52%") { + t.Fatalf("width %d transcript does not contain rendered chart: %q", width, model.transcriptCache) + } + for _, line := range strings.Split(ansi.Strip(model.transcriptCache), "\n") { + if lineWidth := ansi.StringWidth(line); lineWidth > width-1 { + t.Errorf("width %d transcript line width = %d: %q", width, lineWidth, line) + } + } + }) + } +} + +func TestWriteAnswerPlainRendersChartsAndDashboardLinks(t *testing.T) { + sessionID, err := uuid.Parse(testSessionID) + if err != nil { + t.Fatal(err) + } + markdown := metricPieSource + "\n\n[Add to dashboard](#ai-add-to-dashboard:ref)" + var stdout bytes.Buffer + err = writeAnswer(PrintOptions{ + Stdout: &stdout, + APIHost: "api.rootly.com", + SessionID: &sessionID, + IsStdoutTTY: func() bool { return false }, + }, markdown) + if err != nil { + t.Fatal(err) + } + got := stdout.String() + if strings.Contains(got, `{"type"`) || strings.Contains(got, "```metric-chart") { + t.Errorf("plain output contains the raw chart fence: %q", got) + } + if !strings.Contains(got, "[Add to dashboard: open in web](https://rootly.com/account?chat="+testSessionID+")") { + t.Errorf("plain output did not rewrite dashboard link: %q", got) + } + if strings.Contains(got, "\x1b") { + t.Errorf("plain output contains ANSI: %q", got) + } + + var terminalOutput bytes.Buffer + err = writeAnswer(PrintOptions{ + Stdout: &terminalOutput, + APIHost: "api.rootly.com", + SessionID: &sessionID, + IsStdoutTTY: func() bool { return true }, + }, markdown) + if err != nil { + t.Fatal(err) + } + terminalText := ansi.Strip(terminalOutput.String()) + if strings.Contains(terminalText, `{"type"`) || strings.Contains(terminalText, "```metric-chart") || + strings.Contains(strings.ToLower(terminalText), "ai-add-to-dashboard") { + t.Errorf("terminal output contains an unrendered chart or dashboard marker: %q", terminalText) + } + if !strings.Contains(terminalText, "Add to dashboard: open in web") || + !strings.Contains(terminalOutput.String(), "https://rootly.com/account?chat="+testSessionID) { + t.Errorf("terminal output did not render dashboard link: %q", terminalText) + } +} + +func TestExpandStepSeriesKeepsEveryMeasurement(t *testing.T) { + for _, tc := range []struct{ count, width int }{{19, 20}, {7, 60}, {3, 60}, {20, 20}, {90, 20}} { + values := make([]float64, tc.count) + for index := range values { + values[index] = float64(index) + } + expanded := expandStepSeries(values, tc.width) + if len(expanded) != tc.width { + t.Fatalf("%d values at width %d: got %d columns", tc.count, tc.width, len(expanded)) + } + if expanded[0] != 0 || expanded[len(expanded)-1] != float64(tc.count-1) { + t.Errorf("%d values at width %d: first/last = %v/%v", tc.count, tc.width, expanded[0], expanded[len(expanded)-1]) + } + if tc.count <= tc.width { + seen := map[float64]bool{} + for _, value := range expanded { + seen[value] = true + } + if len(seen) != tc.count { + t.Errorf("%d values at width %d: %d distinct measurements plotted", tc.count, tc.width, len(seen)) + } + } + } +} + +func TestLineChartKeepsPeakAtEveryWidth(t *testing.T) { + values := make([]string, 31) + for index := range values { + values[index] = "0" + } + for _, index := range []int{3, 11, 19, 27} { + values[index] = "1" + } + labels := make([]string, len(values)) + for index := range labels { + labels[index] = fmt.Sprintf("%q", fmt.Sprintf("Oct %d", index+1)) + } + source := fmt.Sprintf(`{"type":"line","labels":[%s],"values":[%s]}`, strings.Join(labels, ","), strings.Join(values, ",")) + for _, width := range []int{139, 110, 70, 40} { + lines, ok := renderMetricChart(source, markdownOpts{width: width, noColor: true}) + if !ok { + t.Fatalf("width %d: chart not rendered", width) + } + graph := strings.Join(lines, "\n") + if !strings.Contains(graph, "1.00") || strings.Contains(graph, "0.90") || strings.Contains(graph, "0.79") { + t.Errorf("width %d: peak not preserved:\n%s", width, graph) + } + } +} + +func TestFitLineSeriesBucketsLongSeriesByPeak(t *testing.T) { + values := make([]float64, 90) + values[44] = 7 + series, width := fitLineSeries(values, 30) + if width != 30 || len(series) != 30 { + t.Fatalf("got %d values at width %d, want 30", len(series), width) + } + if peak := slices.Max(series); peak != 7 { + t.Errorf("peak = %v, want 7", peak) + } +} + +func TestFitLineSeriesKeepsExactEndpoints(t *testing.T) { + values := make([]float64, 61) + for index := range values { + values[index] = 7 + } + values[0], values[60] = 2, 0 + values[59] = 9 + series, width := fitLineSeries(values, 60) + if width != 60 || len(series) != 60 { + t.Fatalf("got %d values at width %d, want 60", len(series), width) + } + if series[0] != 2 || series[59] != 0 { + t.Errorf("endpoints = %v/%v, want 2/0", series[0], series[59]) + } + if peak := slices.Max(series); peak != 9 { + t.Errorf("peak = %v, want 9", peak) + } +} diff --git a/internal/agent/polish.go b/internal/agent/polish.go new file mode 100644 index 0000000..4c4f099 --- /dev/null +++ b/internal/agent/polish.go @@ -0,0 +1,164 @@ +package agent + +import ( + "errors" + "fmt" + "strings" + + tea "charm.land/bubbletea/v2" + "charm.land/lipgloss/v2" + "github.com/charmbracelet/x/ansi" + + "github.com/rootlyhq/rootly-cli/internal/mcp" +) + +const ( + sessionNotFoundMessage = "Session not found — starting a new chat" + historyLoadingStatus = "Loading session history…" + blurredCursorGlyph = "▢" + mcpReasonMaxWidth = 120 +) + +var errSessionNotFound = errors.New("session not found") + +type mcpStartupCheckMsg struct{} + +func mcpStartupCheck() tea.Msg { return mcpStartupCheckMsg{} } + +func (m *TUIModel) updateFocus(message tea.Msg) bool { + switch message.(type) { + case tea.FocusMsg: + m.terminalBlurred = false + return true + case tea.BlurMsg: + m.terminalBlurred = true + return true + } + return false +} + +func (m *TUIModel) composerView() string { + if !m.terminalBlurred || !m.composer.Focused() { + return m.composer.View() + } + positioned := m.composer + positioned.SetVirtualCursor(false) + position := positioned.Cursor() + blurred := m.composer + blurred.Blur() + view := blurred.View() + if position == nil { + return view + } + return overlayBlurredCursor(view, position.X, position.Y) +} + +func overlayBlurredCursor(view string, x, y int) string { + glyph := blurredCursorGlyph + lines := strings.Split(view, "\n") + if y < 0 || y >= len(lines) || x < 0 { + return view + } + line := lines[y] + width := ansi.StringWidth(line) + if x >= width { + lines[y] = line + strings.Repeat(" ", x-width) + glyph + } else { + lines[y] = ansi.Cut(line, 0, x) + glyph + ansi.Cut(line, x+1, width) + } + return strings.Join(lines, "\n") +} + +func (m *TUIModel) continueLine() bool { + value := m.composer.Value() + trailing := len(value) - len(strings.TrimRight(value, `\`)) + if trailing%2 == 0 { + return false + } + m.composer.SetValue(strings.TrimSuffix(value, `\`) + "\n") + m.composer.CursorEnd() + m.historyIndex = -1 + return true +} + +func moreLinesHint(hidden int) string { + if hidden == 1 { + return "↑ 1 more line" + } + return fmt.Sprintf("↑ %d more lines", hidden) +} + +func (m *TUIModel) activityLine(line string) string { + line = sanitizeTerminalText(line) + rest, failed := strings.CutPrefix(line, "✗") + if !failed || m.noColor { + return faint(m.noColor, line) + } + return lipgloss.NewStyle().Foreground(lipgloss.Color("1")).Render("✗") + faint(false, rest) +} + +func (m *TUIModel) noteClientToolsAccepted(accepted *int) { + if accepted == nil || m.sentTools == 0 { + return + } + if *accepted == 0 && !m.toolsOffShown { + m.toolsOffShown = true + m.println(faint(m.noColor, localToolsOffMessage)) + return + } + if *accepted > 0 && !m.toolsOnShown { + m.toolsOnShown = true + m.println(faint(m.noColor, localToolsAvailableMessage(*accepted))) + } +} + +func localToolsAvailableMessage(count int) string { + if count == 1 { + return "1 local tool available" + } + return fmt.Sprintf("%d local tools available", count) +} + +func (m *TUIModel) handleHistoryLoaded(msg historyLoadedMsg) { + if errors.Is(msg.err, errSessionNotFound) { + m.startNewSession() + m.println(faint(m.noColor, sessionNotFoundMessage)) + return + } + if msg.err != nil { + m.println(faint(m.noColor, "Could not load session history: "+sanitizeTerminalText(msg.err.Error()))) + return + } + m.printHistory(msg.messages) +} + +func (m *TUIModel) reportUnreachableMCP() { + if m.mcp == nil || m.mcpStartupReported || mcpStillConnecting(m.mcpPage.views) { + return + } + m.mcpStartupReported = true + for _, view := range m.mcpPage.views { + if view.SessionOff || view.Status != mcp.StatusError { + continue + } + m.println(faint(m.noColor, unreachableMCPLine(view))) + } +} + +func mcpStillConnecting(views []mcp.ServerView) bool { + for _, view := range views { + if view.Status == mcp.StatusConnecting && !view.SessionOff { + return true + } + } + return false +} + +func unreachableMCPLine(view mcp.ServerView) string { + reason := strings.Join(strings.Fields(sanitizeTerminalText(view.Detail)), " ") + if reason == "" { + reason = "connection failed" + } + reason = ansi.Truncate(reason, mcpReasonMaxWidth, "…") + return fmt.Sprintf("MCP '%s' unreachable: %s — open /mcp", sanitizeTerminalText(view.Server.Name), reason) +} diff --git a/internal/agent/polish_test.go b/internal/agent/polish_test.go new file mode 100644 index 0000000..011dd43 --- /dev/null +++ b/internal/agent/polish_test.go @@ -0,0 +1,417 @@ +package agent + +import ( + "context" + "errors" + "net/http" + "strings" + "testing" + + tea "charm.land/bubbletea/v2" + "github.com/charmbracelet/x/ansi" + "github.com/google/uuid" + "github.com/oapi-codegen/nullable" + rootly "github.com/rootlyhq/rootly-go" + + "github.com/rootlyhq/rootly-cli/internal/mcp" +) + +func TestFocusReportingSwapsCursorWithoutTouchingComposer(t *testing.T) { + model := NewTUIModel(context.Background(), &fakeStreamer{}, TUIOptions{NoColor: true}) + model.Update(tea.WindowSizeMsg{Width: 80, Height: 24}) + model.Init() + model.composer.SetValue("draft") + + if view := model.View(); !view.ReportFocus { + t.Fatal("View().ReportFocus = false, want focus reporting on") + } + model.Update(tea.BlurMsg{}) + if !model.terminalBlurred { + t.Fatal("terminalBlurred = false after BlurMsg") + } + if !strings.Contains(ansi.Strip(model.composerView()), blurredCursorGlyph) { + t.Errorf("blurred composer = %q, want hollow cursor %q", model.composerView(), blurredCursorGlyph) + } + model.Update(tea.FocusMsg{}) + if model.terminalBlurred { + t.Fatal("terminalBlurred = true after FocusMsg") + } + if strings.Contains(model.composerView(), blurredCursorGlyph) { + t.Errorf("focused composer = %q, want no hollow cursor", model.composerView()) + } + if got := model.composer.Value(); got != "draft" { + t.Errorf("composer value = %q after focus events, want %q", got, "draft") + } + if !model.composer.Focused() { + t.Error("composer lost focus after terminal focus events") + } +} + +func TestOverlayGlyphPlacesCursor(t *testing.T) { + if got := overlayBlurredCursor("ab\ncd", 1, 1); got != "ab\nc▢" { + t.Errorf("overlayBlurredCursor() = %q", got) + } + if got := overlayBlurredCursor("ab", 4, 0); got != "ab ▢" { + t.Errorf("overlayBlurredCursor() past end = %q", got) + } + if got := overlayBlurredCursor("ab", 0, 3); got != "ab" { + t.Errorf("overlayBlurredCursor() out of range = %q", got) + } +} + +func TestTrailingBackslashEnterInsertsNewline(t *testing.T) { + streamer := &fakeStreamer{} + model := NewTUIModel(context.Background(), streamer, TUIOptions{NoColor: true}) + model.Update(tea.WindowSizeMsg{Width: 80, Height: 24}) + model.Init() + model.composer.SetValue(`first line\`) + + model.Update(tea.KeyPressMsg{Code: tea.KeyEnter}) + if got := model.composer.Value(); got != "first line\n" { + t.Fatalf("composer value = %q, want continuation newline", got) + } + if len(streamer.calls) != 0 || strings.Contains(model.Transcript(), "first line") { + t.Fatalf("continuation submitted the prompt (calls=%d)", len(streamer.calls)) + } + + model.composer.SetValue(`path C:\\`) + if model.continueLine() { + t.Error("escaped backslash continued the line, want submit") + } + model.composer.SetValue(`path C:\\\`) + if !model.continueLine() || model.composer.Value() != "path C:\\\\\n" { + t.Errorf("three trailing backslashes = %q, want escaped pair plus newline", model.composer.Value()) + } +} + +func TestComposerOverflowShowsMoreLinesHint(t *testing.T) { + model := NewTUIModel(context.Background(), &fakeStreamer{}, TUIOptions{NoColor: true}) + model.Update(tea.WindowSizeMsg{Width: 80, Height: 12}) + model.Init() + lines := make([]string, 10) + for index := range lines { + lines[index] = "line" + } + model.composer.SetValue(strings.Join(lines, "\n")) + model.composer.CursorEnd() + model.Update(tea.KeyPressMsg{Code: tea.KeyEnd}) + hidden := model.composer.ScrollYOffset() + if hidden == 0 { + t.Fatalf("composer ScrollYOffset = 0 with %d lines and height %d", len(lines), model.composer.Height()) + } + footer := ansi.Strip(model.footerContent()) + if want := moreLinesHint(hidden); !strings.Contains(footer, want) { + t.Errorf("footer = %q, want %q", footer, want) + } + if got := moreLinesHint(1); got != "↑ 1 more line" { + t.Errorf("moreLinesHint(1) = %q", got) + } + if got := moreLinesHint(4); got != "↑ 4 more lines" { + t.Errorf("moreLinesHint(4) = %q", got) + } + + model.composer.SetValue("short") + if footer := ansi.Strip(model.footerContent()); strings.Contains(footer, "more line") { + t.Errorf("footer = %q, want no hint when everything fits", footer) + } +} + +func TestBulletizeWideTablesOnlyWhenWiderThanTerminal(t *testing.T) { + table := "| Service | Owner | Status |\n|---|---|---|\n| api | Platform team | degraded |\n| web | Frontend | ok |" + + if got := bulletizeWideTables(table, 80); got != table { + t.Errorf("narrow table changed:\n%s", got) + } + got := bulletizeWideTables("Summary\n\n"+table+"\n\nDone", 20) + want := "Summary\n\n- **api** · Owner: Platform team · Status: degraded\n- **web** · Owner: Frontend · Status: ok\n\nDone" + if got != want { + t.Errorf("wide table =\n%s\nwant\n%s", got, want) + } + + fenced := "```\n" + table + "\n```" + if got := bulletizeWideTables(fenced, 10); got != fenced { + t.Errorf("fenced table changed:\n%s", got) + } + if got := bulletizeWideTables("a | b\nnot a separator", 2); got != "a | b\nnot a separator" { + t.Errorf("non-table changed: %q", got) + } +} + +func TestBulletizeWideTablesEdgeCases(t *testing.T) { + headerOnly := "| A very long header | Another very long header |\n|---|---|" + if got := bulletizeWideTables(headerOnly, 10); got != headerOnly { + t.Errorf("header-only table changed:\n%s", got) + } + + escaped := "| Rule | Pattern | Owner |\n|---|---|---|\n| split | a \\| b | Platform team |" + want := "- **split** · Pattern: a | b · Owner: Platform team" + if got := bulletizeWideTables(escaped, 10); got != want { + t.Errorf("escaped pipe table =\n%s\nwant\n%s", got, want) + } + + doubled := "| Path | Owner |\n|---|---|\n| C:\\\\ | Platform team |" + wantDoubled := "- **C:\\\\** · Owner: Platform team" + if got := bulletizeWideTables(doubled, 10); got != wantDoubled { + t.Errorf("doubled backslash table =\n%s\nwant\n%s", got, wantDoubled) + } + if got := tableCells(`| a \\| b |`); len(got) != 2 || got[0] != `a \\` || got[1] != "b" { + t.Errorf("tableCells with escaped backslash = %q", got) + } + + table := "| Service | Owner | Status |\n|---|---|---|\n| api | Platform team | degraded |" + longFence := "````md\n```\n" + table + "\n````" + if got := bulletizeWideTables(longFence, 10); got != longFence { + t.Errorf("table inside four-backtick fence changed:\n%s", got) + } +} + +func TestPrepareAssistantMarkdownBulletizesWideTables(t *testing.T) { + table := "| Service | Owner |\n|---|---|\n| api | Platform team |" + prepared, _ := prepareAssistantMarkdown(table, markdownOpts{width: 12}) + if !strings.Contains(prepared, "- **api** · Owner: Platform team") { + t.Errorf("prepared = %q, want bullets", prepared) + } +} + +func TestFailedActivityLineHasRedCross(t *testing.T) { + colored := NewTUIModel(context.Background(), &fakeStreamer{}, TUIOptions{}) + line := colored.activityLine("✗ Searching incidents (1.0s)") + if !strings.Contains(line, "\x1b[31m✗") { + t.Errorf("activity line = %q, want red ✗", line) + } + if ok := colored.activityLine("✓ Searching incidents (1.0s)"); strings.Contains(ok, "\x1b[31m") { + t.Errorf("success line = %q, want no red", ok) + } + plain := NewTUIModel(context.Background(), &fakeStreamer{}, TUIOptions{NoColor: true}) + if got := plain.activityLine("✗ Searching incidents (1.0s)"); got != "✗ Searching incidents (1.0s)" { + t.Errorf("no-color failed line = %q, want plain", got) + } +} + +func TestLocalToolsAvailableStatusShownOnce(t *testing.T) { + model := NewTUIModel(context.Background(), &fakeStreamer{}, TUIOptions{NoColor: true}) + model.sentTools = 3 + accepted := 3 + model.noteClientToolsAccepted(&accepted) + model.noteClientToolsAccepted(&accepted) + transcript := model.Transcript() + if got := strings.Count(transcript, "3 local tools available"); got != 1 { + t.Errorf("status shown %d times, want once:\n%s", got, transcript) + } + if strings.Contains(transcript, localToolsOffMessage) { + t.Errorf("transcript = %q, want no flag-off notice", transcript) + } + + none := NewTUIModel(context.Background(), &fakeStreamer{}, TUIOptions{NoColor: true}) + none.noteClientToolsAccepted(&accepted) + if strings.Contains(none.Transcript(), "local tool") { + t.Error("status shown when no tools were sent") + } + if got := localToolsAvailableMessage(1); got != "1 local tool available" { + t.Errorf("localToolsAvailableMessage(1) = %q", got) + } +} + +func TestResumeMissingSessionStartsFresh(t *testing.T) { + sessionID := uuid.New() + fake := &fakeTUIAPI{messagePages: []*rootly.ListAiChatSessionMessagesResponse{{ + Body: []byte(`{"errors":[{"title":"Not found"}]}`), + HTTPResponse: &http.Response{StatusCode: http.StatusNotFound}, + }}} + model := NewTUIModel(context.Background(), &fakeStreamer{}, TUIOptions{NoColor: true, API: fake, SessionID: &sessionID, SessionTitle: "Old"}) + model.Update(tea.WindowSizeMsg{Width: 80, Height: 24}) + msg := model.loadHistory(sessionID)() + loaded, ok := msg.(historyLoadedMsg) + if !ok || !errors.Is(loaded.err, errSessionNotFound) { + t.Fatalf("loadHistory() = %#v, want errSessionNotFound", msg) + } + model.Update(msg) + if model.sessionID != nil { + t.Errorf("sessionID = %v, want a fresh session", model.sessionID) + } + transcript := model.Transcript() + if !strings.Contains(transcript, sessionNotFoundMessage) { + t.Errorf("transcript = %q, want %q", transcript, sessionNotFoundMessage) + } + if strings.Contains(transcript, "Could not load session history") { + t.Errorf("transcript = %q, want no generic history error", transcript) + } +} + +func TestStartupReportsUnreachableMCPServersOnce(t *testing.T) { + manager := &fakeMCPManager{views: []mcp.ServerView{ + {Server: mcp.Server{Name: "fs"}, Status: mcp.StatusConnected}, + {Server: mcp.Server{Name: "broken"}, Status: mcp.StatusError, Detail: "exec: \"nope\": executable\nfile not found"}, + {Server: mcp.Server{Name: "quiet"}, Status: mcp.StatusError, Detail: "", SessionOff: true}, + {Server: mcp.Server{Name: "blank"}, Status: mcp.StatusError}, + }, changes: make(chan struct{})} + model := NewTUIModel(context.Background(), &fakeStreamer{}, TUIOptions{NoColor: true, MCP: manager}) + model.Update(mcpStartupCheckMsg{}) + model.Update(mcpChangedMsg{}) + + transcript := model.Transcript() + for _, want := range []string{ + `MCP 'broken' unreachable: exec: "nope": executable file not found — open /mcp`, + "MCP 'blank' unreachable: connection failed — open /mcp", + } { + if got := strings.Count(transcript, want); got != 1 { + t.Errorf("transcript has %q %d times, want once:\n%s", want, got, transcript) + } + } + for _, unwanted := range []string{"'fs'", "'quiet'"} { + if strings.Contains(transcript, unwanted) { + t.Errorf("transcript = %q, want no line for %s", transcript, unwanted) + } + } +} + +func TestWelcomeShowsFirstNameWorkingDirAndCapability(t *testing.T) { + user := currentUserForWelcome("Avery Rootly", "avery@example.com", "") + user.Data.Attributes.FirstName = nullable.NewNullableWithValue("Ave") + model := NewTUIModel(context.Background(), &fakeStreamer{}, TUIOptions{NoColor: true, CurrentUser: user, APIHost: "api.rootly.com", WorkingDir: "/srv/ops"}) + model.width = 100 + rows := model.welcomeTextRows() + for _, want := range []string{"Welcome back, Ave", "/srv/ops", "api.rootly.com", welcomeCapabilityText} { + if !strings.Contains(rows, want) { + t.Errorf("welcome rows = %q, want %q", rows, want) + } + } + model.width = 40 + if strings.Contains(model.welcomeTextRows(), welcomeCapabilityText) { + t.Error("capability line shown when it does not fit") + } +} + +func TestPrintExitErrorCodes(t *testing.T) { + var coder interface{ ExitCode() int } + if err := PrintExitError(NotAuthenticatedError("not authenticated")); !errors.As(err, &coder) || coder.ExitCode() != ExitNotAuthenticated { + t.Errorf("PrintExitError(not authenticated) = %v, want exit code 2", err) + } + generic := errors.New("boom") + if err := PrintExitError(generic); err != generic { + t.Errorf("PrintExitError(generic) = %v, want unchanged", err) + } + if err := PrintExitError(nil); err != nil { + t.Errorf("PrintExitError(nil) = %v", err) + } + if !IsNotAuthenticated(StatusResponseError(http.StatusUnauthorized, nil)) { + t.Error("401 stream error is not classified as not authenticated") + } +} + +func TestStaleMissingSessionReplyKeepsNewChat(t *testing.T) { + sessionID := uuid.New() + fake := &fakeTUIAPI{messagePages: []*rootly.ListAiChatSessionMessagesResponse{{ + Body: []byte(`{"errors":[{"title":"Not found"}]}`), + HTTPResponse: &http.Response{StatusCode: http.StatusNotFound}, + }}} + model := NewTUIModel(context.Background(), &fakeStreamer{}, TUIOptions{NoColor: true, API: fake, SessionID: &sessionID, SessionTitle: "Old"}) + model.Update(tea.WindowSizeMsg{Width: 80, Height: 24}) + pending := model.loadHistory(sessionID) + model.startNewSession() + newSession := uuid.New() + model.sessionID = &newSession + model.Update(pending()) + if model.sessionID == nil || *model.sessionID != newSession { + t.Fatalf("sessionID = %v, want the new chat %v kept", model.sessionID, newSession) + } + if strings.Contains(model.Transcript(), sessionNotFoundMessage) { + t.Errorf("stale 404 printed %q", sessionNotFoundMessage) + } +} + +func TestStartupWaitsForConnectingViews(t *testing.T) { + manager := &fakeMCPManager{views: []mcp.ServerView{ + {Server: mcp.Server{Name: "slow"}, Status: mcp.StatusConnecting}, + {Server: mcp.Server{Name: "broken"}, Status: mcp.StatusError, Detail: "boom"}, + }, changes: make(chan struct{})} + model := NewTUIModel(context.Background(), &fakeStreamer{}, TUIOptions{NoColor: true, MCP: manager}) + model.Update(mcpStartupCheckMsg{}) + if strings.Contains(model.Transcript(), "unreachable") { + t.Fatalf("reported before every server settled:\n%s", model.Transcript()) + } + manager.views[0] = mcp.ServerView{Server: mcp.Server{Name: "slow"}, Status: mcp.StatusError, Detail: "timeout"} + model.Update(mcpChangedMsg{}) + transcript := model.Transcript() + for _, want := range []string{"MCP 'slow' unreachable: timeout", "MCP 'broken' unreachable: boom"} { + if strings.Count(transcript, want) != 1 { + t.Errorf("transcript missing %q once:\n%s", want, transcript) + } + } +} + +func TestBulletizeWideTablesAfterLongerFenceCloser(t *testing.T) { + table := "| Service | Owner | Status |\n|---|---|---|\n| api | Platform team | degraded |" + got := bulletizeWideTables("```\ncode\n````\n"+table, 10) + if !strings.Contains(got, "- **api** · Owner: Platform team · Status: degraded") { + t.Errorf("table after a longer fence closer was not converted:\n%s", got) + } +} + +func TestBulletizeWideTablesFenceRules(t *testing.T) { + table := "| Service | Owner | Status |\n|---|---|---|\n| api | Platform team | degraded |" + bullet := "- **api** · Owner: Platform team · Status: degraded" + cases := map[string]struct { + markdown string + convert bool + }{ + "inner shorter fence stays inside": {"````\n```\n" + table + "\n````", false}, + "longer closer closes": {"```\ncode\n````\n" + table, true}, + "indented code is not a fence": {" ```\n\n" + table, true}, + "indented closer does not close": {"```\n ```\n" + table + "\n```", false}, + } + for name, tc := range cases { + t.Run(name, func(t *testing.T) { + got := bulletizeWideTables(tc.markdown, 10) + if strings.Contains(got, bullet) != tc.convert { + t.Errorf("convert=%v, got:\n%s", tc.convert, got) + } + }) + } +} + +func TestBulletizeWideTablesKeepsEscapedPipeInItsColumn(t *testing.T) { + got := bulletizeWideTables("| Service | Status |\n|---|---|\n| api\\|worker | degraded |", 10) + if !strings.Contains(got, "- **api|worker** · Status: degraded") { + t.Errorf("escaped pipe split the row:\n%s", got) + } +} + +func TestSendWaitsForPendingHistory(t *testing.T) { + sessionID := uuid.New() + fake := &fakeTUIAPI{messagePages: []*rootly.ListAiChatSessionMessagesResponse{{ + Body: []byte(`{"errors":[{"title":"Not found"}]}`), + HTTPResponse: &http.Response{StatusCode: http.StatusNotFound}, + }}} + model := NewTUIModel(context.Background(), &fakeStreamer{}, TUIOptions{NoColor: true, API: fake, SessionID: &sessionID}) + model.Update(tea.WindowSizeMsg{Width: 80, Height: 24}) + pending := model.loadHistory(sessionID) + model.composer.SetValue("hello") + if cmd := model.submit(); cmd != nil || model.state == tuiStreaming { + t.Fatal("a prompt sent before resumed history loads must wait") + } + if model.queuedPrompt.message != "hello" || model.status != "Loading session history…" { + t.Fatalf("queued %q status %q", model.queuedPrompt, model.status) + } + _, cmd := model.Update(pending()) + if model.historyPending || model.queuedPrompt != (queuedSend{}) { + t.Fatal("settled history should clear the pending flag and the queue") + } + if cmd == nil || model.state != tuiStreaming || model.currentPrompt != "hello" { + t.Fatal("the queued prompt should send once history has settled") + } +} + +func TestSkillPromptWaitsForPendingHistory(t *testing.T) { + sessionID := uuid.New() + model := NewTUIModel(context.Background(), &fakeStreamer{}, TUIOptions{NoColor: true, API: &fakeTUIAPI{}, SessionID: &sessionID}) + model.Update(tea.WindowSizeMsg{Width: 80, Height: 24}) + model.loadHistory(sessionID) + if cmd := model.submitPrompt("expanded skill prompt"); cmd != nil || model.state == tuiStreaming { + t.Fatal("skill prompts must wait for resumed history too") + } + model.Update(tea.KeyPressMsg{Code: 'c', Mod: tea.ModCtrl}) + if model.queuedPrompt != (queuedSend{}) { + t.Fatal("Ctrl+C should drop a queued prompt") + } +} diff --git a/internal/agent/print.go b/internal/agent/print.go new file mode 100644 index 0000000..1bdcda4 --- /dev/null +++ b/internal/agent/print.go @@ -0,0 +1,533 @@ +package agent + +import ( + "bufio" + "context" + "encoding/json" + "errors" + "fmt" + "io" + "strings" + "time" + + "charm.land/glamour/v2" + "github.com/google/uuid" + rootly "github.com/rootlyhq/rootly-go" + + "github.com/rootlyhq/rootly-cli/internal/api" + "github.com/rootlyhq/rootly-cli/internal/mcp" + "github.com/rootlyhq/rootly-cli/internal/oauth" +) + +type PrintOptions struct { + Prompt string + Stdin io.Reader + Stdout io.Writer + Stderr io.Writer + NoColor bool + Quiet bool + OutputFormat string + APIHost string + SessionID *uuid.UUID + Incident *api.Incident + Alert *api.Alert + Confirmation *rootly.AiChatStreamConfirmationRequired + StartedAt time.Time + IsStdinTTY func() bool + IsStdoutTTY func() bool + IsStderrTTY func() bool + LocalTools LocalTools + ToolResults ToolResultPoster +} + +func RunPrint(ctx context.Context, streamer Streamer, options PrintOptions) error { + options, prompt, err := preparePrint(options) + if err != nil { + return err + } + + runCtx, cancel := context.WithCancel(ctx) + defer cancel() + run := printRun{ + streamer: streamer, + options: options, + params: initialPrintParams(prompt, options), + sessionID: options.SessionID, + ctx: runCtx, + cancel: cancel, + } + return run.run(runCtx) +} + +func preparePrint(options PrintOptions) (PrintOptions, string, error) { + if options.OutputFormat == "" { + options.OutputFormat = "text" + } + if options.OutputFormat != "text" && options.OutputFormat != "json" && options.OutputFormat != "stream-json" { + return options, "", fmt.Errorf("invalid output format %q (expected text, json, or stream-json)", options.OutputFormat) + } + options.StartedAt = time.Now() + prompt, err := readPrompt(options) + if err != nil { + return options, "", err + } + if strings.TrimSpace(prompt) == "" { + return options, "", errors.New("message is required") + } + return options, prompt, nil +} + +func initialPrintParams(prompt string, options PrintOptions) *rootly.StreamAiChatParams { + message := prompt + params := &rootly.StreamAiChatParams{Message: &message, SessionID: options.SessionID} + if options.Incident != nil { + if id, parseErr := uuid.Parse(options.Incident.ID); parseErr == nil { + params.IncidentID = &id + } + } + if options.Alert != nil { + if id, parseErr := uuid.Parse(options.Alert.ID); parseErr == nil { + params.AlertID = &id + } + } + return params +} + +type printRun struct { + streamer Streamer + options PrintOptions + params *rootly.StreamAiChatParams + sessionID *uuid.UUID + answer strings.Builder + activities activityTracker + confirmation *rootly.AiChatStreamConfirmationRequired + doneStatus string + doneReceived bool + streamError error + toolResultErr error + ctx context.Context + cancel context.CancelFunc + sentTools int + toolsOffSeen bool + stdinReader *bufio.Reader + credHintShown bool +} + +func (r *printRun) run(ctx context.Context) error { + for { + continues, err := r.runTurn(ctx) + if err != nil || !continues { + return err + } + } +} + +func (r *printRun) runTurn(ctx context.Context) (bool, error) { + r.confirmation = nil + r.doneStatus = "" + r.doneReceived = false + r.streamError = nil + r.toolResultErr = nil + resumeTextSeparator := r.params.CheckpointID != nil + r.ctx = ctx + r.sentTools = 0 + streamCtx := ctx + if r.options.LocalTools != nil { + tools, _ := r.options.LocalTools.ClientTools() + r.sentTools = len(tools) + streamCtx = api.WithClientTools(ctx, tools) + } + err := StreamTurn(streamCtx, r.streamer, r.params, func(event Event) bool { + return r.handleEvent(event, &resumeTextSeparator) + }) + if r.toolResultErr != nil { + r.writeActivityLines(r.activities.Finish(false)) + return false, finishPrintError(r.options, r.answer.String(), r.toolResultErr) + } + if err != nil { + r.writeActivityLines(r.activities.Finish(false)) + return false, finishPrintError(r.options, r.answer.String(), err) + } + if r.streamError != nil { + return false, finishPrintError(r.options, r.answer.String(), r.streamError) + } + if !r.doneReceived { + return false, finishPrintError(r.options, r.answer.String(), errors.New("AI chat stream ended before the turn finished")) + } + if r.doneStatus == "error" { + return false, finishPrintError(r.options, r.answer.String(), errors.New("AI chat turn failed")) + } + if r.confirmation != nil || r.doneStatus == "confirmation_required" { + r.writeActivityLines(r.activities.Finish(true)) + return r.resumeConfirmation(r.confirmation) + } + r.options.SessionID = r.sessionID + return false, finishPrintTurn(r.options, r.answer.String(), r.doneStatus) +} + +func (r *printRun) handleEvent(event Event, resumeTextSeparator *bool) bool { + r.writeActivityLines(r.activities.Handle(event)) + if r.options.OutputFormat == "stream-json" { + writeStreamEvent(r.options.Stdout, event) + } + switch event.Type { + case rootly.AiChatStreamEventSessionID: + if event.SessionID != nil { + id, parseErr := uuid.Parse(event.SessionID.SessionID) + if parseErr == nil { + r.sessionID = &id + r.options.SessionID = &id + } + } + if event.ClientToolsAccepted != nil && *event.ClientToolsAccepted == 0 && r.sentTools > 0 && !r.toolsOffSeen { + r.toolsOffSeen = true + if !r.options.Quiet { + _, _ = fmt.Fprintln(r.options.Stderr, faint(r.options.NoColor, localToolsOffMessage)) + } + } + case EventClientToolCall: + r.handleClientToolCall(event.ClientToolCall) + if r.toolResultErr != nil { + return false + } + case rootly.AiChatStreamEventText: + if event.Text != nil { + if *resumeTextSeparator { + if r.answer.Len() > 0 && !strings.HasSuffix(r.answer.String(), "\n") { + r.answer.WriteString("\n\n") + } + *resumeTextSeparator = false + } + content := event.Text.Content + if r.options.OutputFormat == "text" { + content = sanitizeTerminalText(content) + } + r.answer.WriteString(content) + } + case rootly.AiChatStreamEventStatus: + if event.Status != nil { + writeProgress(r.options, sanitizeTerminalText(event.Status.Status)) + } + case rootly.AiChatStreamEventTaskUpdate: + if event.TaskUpdate != nil { + writeProgress(r.options, fmt.Sprintf("task: %s (%s)", sanitizeTerminalText(event.TaskUpdate.Title), sanitizeTerminalText(event.TaskUpdate.Status))) + } + case rootly.AiChatStreamEventSkillLoaded: + if event.SkillLoaded != nil { + skillName := sanitizeTerminalText(event.SkillLoaded.Skill.Name) + writeProgress(r.options, faint(r.options.NoColor, "• Loaded skill "+skillName)) + } + case rootly.AiChatStreamEventConfirmationRequired: + r.confirmation = event.ConfirmationRequired + case rootly.AiChatStreamEventError: + if event.Error != nil { + r.streamError = errors.New(event.Error.Message) + return false + } + case rootly.AiChatStreamEventDone: + r.doneReceived = true + if event.Done != nil { + r.doneStatus = event.Done.Status + } + } + return true +} + +func (r *printRun) handleClientToolCall(call *ClientToolCall) { + if call == nil { + return + } + activity := r.activities.pending + ctx := r.ctx + if ctx == nil { + ctx = context.Background() + } + ref, known := resolveLocalCall(r.options.LocalTools, call) + result := mcp.CallResult{Output: "Unknown local tool " + call.Name, IsError: true} + if known { + writeProgress(r.options, faint(r.options.NoColor, localToolActivity(ref))) + if !r.options.LocalTools.NeedsApproval(call.Name, call.Arguments) || r.approveLocalTool(ref, call) { + result = runLocalCall(ctx, r.options.LocalTools, call) + } else { + result = mcp.CallResult{Output: userDeclinedOutput, IsError: true} + } + } + if result.IsError { + r.activities.discardPendingIf(activity) + } + if looksLikeExpiredCreds(result) && !r.credHintShown && !r.options.Quiet { + r.credHintShown = true + _, _ = fmt.Fprintln(r.options.Stderr, localCredsHint) + } + sessionID := "" + if r.sessionID != nil { + sessionID = r.sessionID.String() + } + if err := postLocalResult(ctx, r.options.ToolResults, sessionID, call, result); err != nil { + r.toolResultErr = sanitizeTerminalError(err) + _, _ = fmt.Fprintln(r.options.Stderr, r.toolResultErr) + if r.cancel != nil { + r.cancel() + } + } +} + +// approveLocalTool asks y/n for a non-read-only tool; without a terminal it declines. +func (r *printRun) approveLocalTool(ref mcp.ToolRef, call *ClientToolCall) bool { + if !terminal(r.options.IsStdinTTY) || !terminal(r.options.IsStderrTTY) || r.options.Stdin == nil { + if !r.options.Quiet { + _, _ = fmt.Fprintf(r.options.Stderr, "Declined local tool %s · %s (approval needs a terminal)\n", sanitizeTerminalText(ref.ServerName), sanitizeTerminalText(ref.Tool)) + } + return false + } + _, _ = fmt.Fprintf(r.options.Stderr, "%s\n", renderLocalToolCard(ref, call, mcp.ApprovalOptions{}, 80, r.options.NoColor)) + _, _ = fmt.Fprint(r.options.Stderr, "Run this local tool? [y/N] ") + if r.stdinReader == nil { + r.stdinReader = bufio.NewReader(r.options.Stdin) + } + line, _ := r.stdinReader.ReadString('\n') + return strings.EqualFold(strings.TrimSpace(line), "y") +} + +func (r *printRun) writeActivityLines(lines []string) { + for _, line := range lines { + writeProgress(r.options, faint(r.options.NoColor, sanitizeTerminalText(line))) + } +} + +func (r *printRun) resumeConfirmation(confirmation *rootly.AiChatStreamConfirmationRequired) (bool, error) { + if confirmation == nil { + err := errors.New("action requires approval (run \"rootly agent\" interactively)") + return false, finishPrintError(r.options, r.answer.String(), err) + } + if !terminal(r.options.IsStdinTTY) || !terminal(r.options.IsStderrTTY) { + if r.options.OutputFormat != "text" { + r.options.Confirmation = confirmation + _ = finishPrintTurn(r.options, r.answer.String(), "confirmation_required") + } + err := fmt.Errorf("action requires approval: %s (run \"rootly agent\" interactively)", sanitizeTerminalText(confirmation.Question)) + if r.options.OutputFormat == "text" { + err = sanitizeTerminalError(err) + } + return false, err + } + decision, err := askForApproval(r.options, confirmation) + if err != nil { + return false, err + } + checkpointID, err := uuid.Parse(confirmation.CheckpointID) + if err != nil { + return false, fmt.Errorf("invalid confirmation checkpoint: %w", err) + } + if r.sessionID == nil { + return false, errors.New("confirmation checkpoint did not include a session id") + } + r.options.SessionID = r.sessionID + if expired(confirmation) { + return false, finishPrintTurn(r.options, r.answer.String(), "expired") + } + r.params = &rootly.StreamAiChatParams{ + SessionID: r.sessionID, + CheckpointID: &checkpointID, + Decision: &decision, + } + return true, nil +} + +func finishPrintError(options PrintOptions, answer string, err error) error { + if options.OutputFormat != "text" { + if outputErr := finishPrintTurn(options, answer, "error"); outputErr != nil { + return outputErr + } + return sanitizeTerminalError(err) + } + return sanitizeTerminalError(err) +} + +func finishPrintTurn(options PrintOptions, answer, doneStatus string) error { + if options.OutputFormat == "json" || options.OutputFormat == "stream-json" { + if doneStatus == "" { + doneStatus = "complete" + } + session := "" + if options.SessionID != nil { + session = options.SessionID.String() + } + result := map[string]any{ + "session_id": session, + "result": answer, + "status": doneStatus, + "duration_ms": time.Since(options.StartedAt).Milliseconds(), + } + if options.OutputFormat == "stream-json" { + result["type"] = "result" + } + if options.Confirmation != nil { + result["confirmation"] = map[string]any{ + "question": options.Confirmation.Question, + "checkpoint_id": options.Confirmation.CheckpointID, + "preview": options.Confirmation.Preview, + "expires_at": options.Confirmation.ExpiresAt, + } + } + encoder := json.NewEncoder(options.Stdout) + return encoder.Encode(result) + } + answer = sanitizeTerminalText(answer) + doneStatus = sanitizeTerminalText(doneStatus) + if strings.TrimSpace(answer) != "" { + if err := writeAnswer(options, answer); err != nil { + return err + } + } + if options.SessionID != nil && terminal(options.IsStderrTTY) { + _, _ = fmt.Fprintf(options.Stderr, "Resume with: rootly agent --resume %s\n", options.SessionID.String()) + } + switch doneStatus { + case "cancelled": + _, err := fmt.Fprintln(options.Stderr, "Action cancelled.") + return err + case "expired": + _, err := fmt.Fprintln(options.Stderr, "Confirmation expired; no action was taken.") + return err + } + if doneStatus != "" && doneStatus != "complete" && doneStatus != "confirmed" { + _, err := fmt.Fprintf(options.Stderr, "turn halted: %s\n", doneStatus) + return err + } + return nil +} + +func readPrompt(options PrintOptions) (string, error) { + prompt := options.Prompt + if options.Stdin == nil || terminal(options.IsStdinTTY) { + return prompt, nil + } + input, err := io.ReadAll(options.Stdin) + if err != nil { + return "", fmt.Errorf("read stdin: %w", err) + } + stdinText := strings.Trim(string(input), "\r\n") + if stdinText == "" { + return prompt, nil + } + if prompt == "" { + return stdinText, nil + } + return strings.TrimRight(prompt, "\r\n") + "\n\n" + stdinText, nil +} + +func writeProgress(options PrintOptions, line string) { + if options.OutputFormat == "text" && !options.Quiet && terminal(options.IsStderrTTY) { + _, _ = fmt.Fprintln(options.Stderr, line) + } +} + +func writeStreamEvent(writer io.Writer, event Event) { + payload := any(nil) + switch event.Type { + case rootly.AiChatStreamEventSessionID: + payload = event.SessionID + case EventClientToolCall: + payload = event.ClientToolCall + case rootly.AiChatStreamEventText: + payload = event.Text + case rootly.AiChatStreamEventStatus: + payload = event.Status + case rootly.AiChatStreamEventTaskUpdate: + payload = event.TaskUpdate + case rootly.AiChatStreamEventSkillLoaded: + payload = event.SkillLoaded + case rootly.AiChatStreamEventConfirmationRequired: + payload = event.ConfirmationRequired + case rootly.AiChatStreamEventCheckpointResolved: + payload = event.CheckpointResolved + case rootly.AiChatStreamEventError: + payload = event.Error + case rootly.AiChatStreamEventDone: + payload = event.Done + } + fields := map[string]any{"type": event.Type} + if payload != nil { + encoded, err := json.Marshal(payload) + if err == nil { + _ = json.Unmarshal(encoded, &fields) + fields["type"] = event.Type + } + } + _ = json.NewEncoder(writer).Encode(fields) +} + +func askForApproval(options PrintOptions, confirmation *rootly.AiChatStreamConfirmationRequired) (rootly.StreamAiChatParamsDecision, error) { + _, _ = fmt.Fprintf(options.Stderr, "%s\n", renderConfirmationCard(confirmation, 80, options.NoColor, time.Now())) + _, _ = fmt.Fprint(options.Stderr, "Approve? [y/N] ") + line, err := bufio.NewReader(options.Stdin).ReadString('\n') + if err != nil && !errors.Is(err, io.EOF) { + return "", fmt.Errorf("read approval: %w", err) + } + if strings.EqualFold(strings.TrimSpace(line), "y") { + return rootly.StreamAiChatParamsDecisionConfirm, nil + } + return rootly.StreamAiChatParamsDecisionCancel, nil +} + +func writeAnswer(options PrintOptions, answer string) error { + answer = sanitizeTerminalText(answer) + sessionID := "" + if options.SessionID != nil { + sessionID = options.SessionID.String() + } + useGlamour := terminal(options.IsStdoutTTY) && !options.NoColor + webBase := "" + if options.APIHost != "" { + webBase = oauth.DeriveAuthBaseURL(options.APIHost) + } + markdownOptions := markdownOpts{ + width: 80, + noColor: !useGlamour, + webBase: webBase, + sessionID: sessionID, + } + if useGlamour { + source, charts := prepareAssistantMarkdown(answer, markdownOptions) + renderer, err := glamour.NewTermRenderer(glamour.WithEnvironmentConfig(), glamour.WithWordWrap(80)) + if err != nil { + return err + } + source, links := preprocessMarkdownLinksWithURLs(source, options.APIHost) + rendered, err := renderer.Render(source) + if err != nil { + return err + } + answer = expandAssistantMarkdownCharts(strings.TrimSpace(restoreMarkdownLinks(rendered, links)), charts) + } else { + answer = renderAssistantMarkdownPlain(answer, markdownOptions) + } + _, err := fmt.Fprintln(options.Stdout, answer) + return err +} + +func renderAssistantMarkdownPlain(markdown string, opts markdownOpts) string { + prepared, charts := prepareAssistantMarkdown(markdown, opts) + return expandAssistantMarkdownCharts(prepared, charts) +} + +func terminal(check func() bool) bool { + return check != nil && check() +} + +func indent(value, prefix string) string { + lines := strings.Split(value, "\n") + for index, line := range lines { + lines[index] = prefix + line + } + return strings.Join(lines, "\n") +} + +func expired(confirmation *rootly.AiChatStreamConfirmationRequired) bool { + if confirmation.ExpiresAt == "" { + return false + } + expiresAt, err := time.Parse(time.RFC3339, confirmation.ExpiresAt) + return err == nil && time.Now().After(expiresAt) +} diff --git a/internal/agent/print_test.go b/internal/agent/print_test.go new file mode 100644 index 0000000..59a60f6 --- /dev/null +++ b/internal/agent/print_test.go @@ -0,0 +1,478 @@ +package agent + +import ( + "bytes" + "context" + "encoding/json" + "errors" + "iter" + "strings" + "testing" + "time" + + rootly "github.com/rootlyhq/rootly-go" + + "github.com/rootlyhq/rootly-cli/internal/api" +) + +const testSessionID = "39fd7596-152c-43d3-afce-48b792ae386d" +const testCheckpointID = "8a348996-dfea-4c76-8c50-cfd0f31f8e5f" + +func TestRunPrintJoinsTextToStdout(t *testing.T) { + streamer := &fakeStreamer{responses: [][]rootly.AiChatStreamEvent{{ + eventOf(rootly.AiChatStreamEventText, `{"content":"Hello "}`), + eventOf(rootly.AiChatStreamEventText, `{"content":"there."}`), + eventOf(rootly.AiChatStreamEventDone, `{"status":"complete"}`), + }}} + var stdout, stderr bytes.Buffer + options := newPrintOptions(&stdout, &stderr) + + if err := RunPrint(context.Background(), streamer, options); err != nil { + t.Fatalf("RunPrint() error = %v", err) + } + if got := stdout.String(); got != "Hello there.\n" { + t.Errorf("stdout = %q, want joined response", got) + } +} + +type blockingToolResultStreamer struct { + stopped chan struct{} +} + +func (s *blockingToolResultStreamer) StreamAiChatEvents(ctx context.Context, _ *rootly.StreamAiChatParams) (iter.Seq2[rootly.AiChatStreamEvent, error], error) { + return func(yield func(rootly.AiChatStreamEvent, error) bool) { + yield(eventOf(EventClientToolCall, `{"call_id":"tool-1","name":"local_unknown","arguments":{}}`), nil) + <-ctx.Done() + close(s.stopped) + }, nil +} + +type failingToolResultPoster struct { + err error +} + +func (p failingToolResultPoster) PostAiChatToolResult(context.Context, api.ToolResult) error { + return p.err +} + +func TestRunPrintCancelsStreamWhenToolResultPostFails(t *testing.T) { + streamer := &blockingToolResultStreamer{stopped: make(chan struct{})} + var stdout, stderr bytes.Buffer + options := newPrintOptions(&stdout, &stderr) + options.Quiet = true + options.ToolResults = failingToolResultPoster{err: errors.New("post failed \x1b]52;c;bad\u0007")} + + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + done := make(chan error, 1) + go func() { + done <- RunPrint(ctx, streamer, options) + }() + select { + case err := <-done: + if err == nil || !strings.Contains(err.Error(), "post failed") { + t.Fatalf("RunPrint() error = %v, want tool-result post error", err) + } + case <-time.After(2 * time.Second): + cancel() + t.Fatal("RunPrint() did not return promptly after tool-result posting failed") + } + if !strings.Contains(stderr.String(), "post failed") { + t.Errorf("stderr = %q, want the tool-result post error even in quiet mode", stderr.String()) + } + if strings.Contains(stderr.String(), "\x1b") || strings.Contains(stderr.String(), "]52;") { + t.Errorf("stderr = %q, want terminal controls sanitized", stderr.String()) + } + select { + case <-streamer.stopped: + default: + t.Fatal("stream was not canceled after tool-result posting failed") + } +} + +func TestRunPrintTextSanitizesAnswerAndStatus(t *testing.T) { + streamer := &fakeStreamer{responses: [][]rootly.AiChatStreamEvent{{ + eventOf(rootly.AiChatStreamEventStatus, `{"status":"\u001b]52;c;ZXZpbA==\u0007Status"}`), + eventOf(rootly.AiChatStreamEventText, `{"content":"Answer \u001b]52;c;ZXZpbA==\u0007 \u001b[2J \u001b]8;;https://evil.example\u0007"}`), + eventOf(rootly.AiChatStreamEventDone, `{"status":"complete"}`), + }}} + var stdout, stderr bytes.Buffer + options := newPrintOptions(&stdout, &stderr) + options.IsStderrTTY = func() bool { return true } + + if err := RunPrint(context.Background(), streamer, options); err != nil { + t.Fatalf("RunPrint() error = %v", err) + } + for label, output := range map[string]string{"stdout": stdout.String(), "stderr": stderr.String()} { + for _, unwanted := range []string{"]52;", "[2J", "\x1b]8;;"} { + if strings.Contains(output, unwanted) { + t.Errorf("%s = %q, want no %q", label, output, unwanted) + } + } + } + if !strings.Contains(stdout.String(), "Answer") || !strings.Contains(stderr.String(), "Status") { + t.Errorf("stdout/stderr = %q / %q, want sanitized answer and status", stdout.String(), stderr.String()) + } +} + +func TestRunPrintJSONAndStreamJSON(t *testing.T) { + rawAnswer := metricPieSource + "\n\n[Add to dashboard](#ai-add-to-dashboard:ref)" + textPayload, err := json.Marshal(map[string]string{"content": rawAnswer}) + if err != nil { + t.Fatal(err) + } + for _, outputFormat := range []string{"json", "stream-json"} { + t.Run(outputFormat, func(t *testing.T) { + streamer := &fakeStreamer{responses: [][]rootly.AiChatStreamEvent{{ + eventOf(rootly.AiChatStreamEventText, string(textPayload)), + eventOf(rootly.AiChatStreamEventDone, `{"status":"complete"}`), + }}} + var stdout, stderr bytes.Buffer + options := newPrintOptions(&stdout, &stderr) + options.OutputFormat = outputFormat + if err := RunPrint(context.Background(), streamer, options); err != nil { + t.Fatalf("RunPrint() error = %v", err) + } + lines := strings.Split(strings.TrimSpace(stdout.String()), "\n") + wantLines := 1 + if outputFormat == "stream-json" { + wantLines = 3 + var event map[string]any + if err := json.Unmarshal([]byte(lines[0]), &event); err != nil { + t.Fatalf("event line is invalid JSON: %v", err) + } + if event["type"] != string(rootly.AiChatStreamEventText) { + t.Errorf("event type = %v, want %q", event["type"], rootly.AiChatStreamEventText) + } + } + if len(lines) != wantLines { + t.Fatalf("stdout lines = %d, want %d: %q", len(lines), wantLines, stdout.String()) + } + var result map[string]any + if err := json.Unmarshal([]byte(lines[len(lines)-1]), &result); err != nil { + t.Fatalf("result line is invalid JSON: %v", err) + } + if outputFormat == "stream-json" && result["type"] != "result" { + t.Errorf("final type = %v, want result", result["type"]) + } + if result["result"] != rawAnswer || result["status"] != "complete" { + t.Errorf("result = %#v, want raw markdown and complete status", result) + } + if stderr.Len() != 0 { + t.Errorf("stderr = %q, want no progress output", stderr.String()) + } + }) + } +} + +func TestRunPrintJSONConfirmationRequiredIsEmittedBeforeError(t *testing.T) { + streamer := &fakeStreamer{responses: [][]rootly.AiChatStreamEvent{{ + eventOf(rootly.AiChatStreamEventSessionID, `{"session_id":"`+testSessionID+`"}`), + eventOf(rootly.AiChatStreamEventConfirmationRequired, `{"checkpoint_id":"`+testCheckpointID+`","question":"Approve?","preview":{"title":"Create"},"expires_at":"2099-01-01T00:00:00Z"}`), + eventOf(rootly.AiChatStreamEventDone, `{"status":"confirmation_required"}`), + }}} + var stdout, stderr bytes.Buffer + options := newPrintOptions(&stdout, &stderr) + options.OutputFormat = "json" + + if err := RunPrint(context.Background(), streamer, options); err == nil { + t.Fatal("RunPrint() returned nil, want non-interactive confirmation error") + } + var result map[string]any + if err := json.Unmarshal(stdout.Bytes(), &result); err != nil { + t.Fatalf("stdout is not a JSON result: %v (%q)", err, stdout.String()) + } + confirmation, ok := result["confirmation"].(map[string]any) + if result["status"] != "confirmation_required" || !ok || confirmation["question"] != "Approve?" { + t.Errorf("result = %#v, want confirmation_required with question", result) + } + if stderr.Len() != 0 { + t.Errorf("stderr = %q, want no JSON-mode progress", stderr.String()) + } +} + +func TestRunPrintReturnsErrorEvent(t *testing.T) { + streamer := &fakeStreamer{responses: [][]rootly.AiChatStreamEvent{{ + eventOf(rootly.AiChatStreamEventError, `{"message":"the turn failed"}`), + eventOf(rootly.AiChatStreamEventDone, `{"status":"error"}`), + }}} + var stdout, stderr bytes.Buffer + err := RunPrint(context.Background(), streamer, newPrintOptions(&stdout, &stderr)) + if err == nil || err.Error() != "the turn failed" { + t.Fatalf("RunPrint() error = %v, want the event error", err) + } + if stdout.Len() != 0 { + t.Errorf("stdout = %q, want no final answer", stdout.String()) + } +} + +func TestRunPrintTreatsMissingDoneAsError(t *testing.T) { + for _, outputFormat := range []string{"text", "json"} { + t.Run(outputFormat, func(t *testing.T) { + streamer := &fakeStreamer{responses: [][]rootly.AiChatStreamEvent{{ + eventOf(rootly.AiChatStreamEventText, `{"content":"Partial answer"}`), + }}} + var stdout, stderr bytes.Buffer + options := newPrintOptions(&stdout, &stderr) + options.OutputFormat = outputFormat + + err := RunPrint(context.Background(), streamer, options) + if err == nil || err.Error() != "AI chat stream ended before the turn finished" { + t.Fatalf("RunPrint() error = %v, want truncated-stream error", err) + } + if outputFormat == "json" { + var result map[string]any + if err := json.Unmarshal(stdout.Bytes(), &result); err != nil { + t.Fatalf("stdout is not a JSON result: %v (%q)", err, stdout.String()) + } + if result["status"] != "error" || result["result"] != "Partial answer" { + t.Errorf("result = %#v, want error status and retained partial answer", result) + } + } + }) + } +} + +func TestRunPrintHaltedStatusPrintsAnswerAndNote(t *testing.T) { + streamer := &fakeStreamer{responses: [][]rootly.AiChatStreamEvent{{ + eventOf(rootly.AiChatStreamEventText, `{"content":"Partial answer"}`), + eventOf(rootly.AiChatStreamEventDone, `{"status":"interrupted"}`), + }}} + var stdout, stderr bytes.Buffer + if err := RunPrint(context.Background(), streamer, newPrintOptions(&stdout, &stderr)); err != nil { + t.Fatalf("RunPrint() error = %v", err) + } + if got := stdout.String(); got != "Partial answer\n" { + t.Errorf("stdout = %q, want partial answer", got) + } + if got := stderr.String(); got != "turn halted: interrupted\n" { + t.Errorf("stderr = %q, want halt note", got) + } +} + +func TestRunPrintConfirmationResumesWithConfirm(t *testing.T) { + streamer := &fakeStreamer{responses: [][]rootly.AiChatStreamEvent{ + { + eventOf(rootly.AiChatStreamEventSessionID, `{"session_id":"`+testSessionID+`"}`), + eventOf(rootly.AiChatStreamEventConfirmationRequired, `{"checkpoint_id":"`+testCheckpointID+`","question":"Create the service?","preview":{"name":"api"},"expires_at":"2099-01-01T00:00:00Z"}`), + eventOf(rootly.AiChatStreamEventDone, `{"status":"confirmation_required"}`), + }, + { + eventOf(rootly.AiChatStreamEventText, `{"content":"Created."}`), + eventOf(rootly.AiChatStreamEventDone, `{"status":"complete"}`), + }, + }} + var stdout, stderr bytes.Buffer + options := newPrintOptions(&stdout, &stderr) + options.Stdin = strings.NewReader("y\n") + options.IsStdinTTY = func() bool { return true } + options.IsStderrTTY = func() bool { return true } + + if err := RunPrint(context.Background(), streamer, options); err != nil { + t.Fatalf("RunPrint() error = %v", err) + } + if got := stdout.String(); got != "Created.\n" { + t.Errorf("stdout = %q, want replayed response", got) + } + if len(streamer.calls) != 2 { + t.Fatalf("stream calls = %d, want 2", len(streamer.calls)) + } + resume := streamer.calls[1] + if resume.Message != nil { + t.Errorf("resume message = %v, want nil", resume.Message) + } + if resume.SessionID == nil || resume.SessionID.String() != testSessionID { + t.Errorf("resume session = %v, want %s", resume.SessionID, testSessionID) + } + if resume.CheckpointID == nil || resume.CheckpointID.String() != testCheckpointID { + t.Errorf("resume checkpoint = %v, want %s", resume.CheckpointID, testCheckpointID) + } + if resume.Decision == nil || *resume.Decision != rootly.StreamAiChatParamsDecisionConfirm { + t.Errorf("resume decision = %v, want confirm", resume.Decision) + } + if !strings.Contains(stderr.String(), "Approve? [y/N]") || !strings.Contains(stderr.String(), `"name": "api"`) { + t.Errorf("stderr = %q, want question prompt and indented preview", stderr.String()) + } +} + +func TestRunPrintConfirmationSeparatesResumedText(t *testing.T) { + streamer := &fakeStreamer{responses: [][]rootly.AiChatStreamEvent{ + { + eventOf(rootly.AiChatStreamEventSessionID, `{"session_id":"`+testSessionID+`"}`), + eventOf(rootly.AiChatStreamEventText, `{"content":"A"}`), + eventOf(rootly.AiChatStreamEventConfirmationRequired, `{"checkpoint_id":"`+testCheckpointID+`","question":"Continue?","expires_at":"2099-01-01T00:00:00Z"}`), + eventOf(rootly.AiChatStreamEventDone, `{"status":"confirmation_required"}`), + }, + { + eventOf(rootly.AiChatStreamEventText, `{"content":"B"}`), + eventOf(rootly.AiChatStreamEventDone, `{"status":"complete"}`), + }, + }} + var stdout, stderr bytes.Buffer + options := newPrintOptions(&stdout, &stderr) + options.Stdin = strings.NewReader("y\n") + options.IsStdinTTY = func() bool { return true } + options.IsStderrTTY = func() bool { return true } + + if err := RunPrint(context.Background(), streamer, options); err != nil { + t.Fatalf("RunPrint() error = %v", err) + } + if got := stdout.String(); got != "A\n\nB\n" { + t.Errorf("stdout = %q, want resumed text separated by a blank line", got) + } +} + +func TestRunPrintConfirmationResumesWithCancel(t *testing.T) { + streamer := &fakeStreamer{responses: [][]rootly.AiChatStreamEvent{ + { + eventOf(rootly.AiChatStreamEventSessionID, `{"session_id":"`+testSessionID+`"}`), + eventOf(rootly.AiChatStreamEventConfirmationRequired, `{"checkpoint_id":"`+testCheckpointID+`","question":"Delete the service?","expires_at":"2099-01-01T00:00:00Z"}`), + eventOf(rootly.AiChatStreamEventDone, `{"status":"confirmation_required"}`), + }, + {eventOf(rootly.AiChatStreamEventDone, `{"status":"cancelled"}`)}, + }} + var stdout, stderr bytes.Buffer + options := newPrintOptions(&stdout, &stderr) + options.Stdin = strings.NewReader("n\n") + options.IsStdinTTY = func() bool { return true } + options.IsStderrTTY = func() bool { return true } + + if err := RunPrint(context.Background(), streamer, options); err != nil { + t.Fatalf("RunPrint() error = %v", err) + } + if got := streamer.calls[1].Decision; got == nil || *got != rootly.StreamAiChatParamsDecisionCancel { + t.Errorf("resume decision = %v, want cancel", got) + } + if got := stdout.String(); got != "" { + t.Errorf("stdout = %q, want no answer for cancellation without replay", got) + } + if !strings.Contains(stderr.String(), "Action cancelled.") || strings.Contains(stderr.String(), "turn halted") { + t.Errorf("stderr = %q, want cancellation note without halted status", stderr.String()) + } +} + +func TestRunPrintExpiredConfirmationDoesNotResume(t *testing.T) { + streamer := &fakeStreamer{responses: [][]rootly.AiChatStreamEvent{{ + eventOf(rootly.AiChatStreamEventSessionID, `{"session_id":"`+testSessionID+`"}`), + eventOf(rootly.AiChatStreamEventConfirmationRequired, `{"checkpoint_id":"`+testCheckpointID+`","question":"Delete the service?","expires_at":"2000-01-01T00:00:00Z"}`), + eventOf(rootly.AiChatStreamEventDone, `{"status":"confirmation_required"}`), + }}} + var stdout, stderr bytes.Buffer + options := newPrintOptions(&stdout, &stderr) + options.Stdin = strings.NewReader("y\n") + options.IsStdinTTY = func() bool { return true } + options.IsStderrTTY = func() bool { return true } + + if err := RunPrint(context.Background(), streamer, options); err != nil { + t.Fatalf("RunPrint() error = %v", err) + } + if got := len(streamer.calls); got != 1 { + t.Errorf("stream calls = %d, want only the original request", got) + } + if !strings.Contains(stderr.String(), "Confirmation expired; no action was taken.") { + t.Errorf("stderr = %q, want expiration message", stderr.String()) + } +} + +func TestRunPrintConfirmationResumesWithConfirmedStatus(t *testing.T) { + streamer := &fakeStreamer{responses: [][]rootly.AiChatStreamEvent{ + { + eventOf(rootly.AiChatStreamEventSessionID, `{"session_id":"`+testSessionID+`"}`), + eventOf(rootly.AiChatStreamEventConfirmationRequired, `{"checkpoint_id":"`+testCheckpointID+`","question":"Create the service?","expires_at":"2099-01-01T00:00:00Z"}`), + eventOf(rootly.AiChatStreamEventDone, `{"status":"confirmation_required"}`), + }, + { + eventOf(rootly.AiChatStreamEventText, `{"content":"Created."}`), + eventOf(rootly.AiChatStreamEventDone, `{"status":"confirmed"}`), + }, + }} + var stdout, stderr bytes.Buffer + options := newPrintOptions(&stdout, &stderr) + options.Stdin = strings.NewReader("y\n") + options.IsStdinTTY = func() bool { return true } + options.IsStderrTTY = func() bool { return true } + + if err := RunPrint(context.Background(), streamer, options); err != nil { + t.Fatalf("RunPrint() error = %v", err) + } + if got := stdout.String(); got != "Created.\n" { + t.Errorf("stdout = %q, want completed answer", got) + } + if strings.Contains(stderr.String(), "turn halted") { + t.Errorf("stderr = %q, confirmed should be treated as complete", stderr.String()) + } +} + +func TestRunPrintExpiredStatus(t *testing.T) { + streamer := &fakeStreamer{responses: [][]rootly.AiChatStreamEvent{{ + eventOf(rootly.AiChatStreamEventDone, `{"status":"expired"}`), + }}} + var stdout, stderr bytes.Buffer + if err := RunPrint(context.Background(), streamer, newPrintOptions(&stdout, &stderr)); err != nil { + t.Fatalf("RunPrint() error = %v", err) + } + if got := stdout.String(); got != "" { + t.Errorf("stdout = %q, want no answer", got) + } + if got := stderr.String(); got != "Confirmation expired; no action was taken.\n" { + t.Errorf("stderr = %q, want expiration message", got) + } +} + +func TestRunPrintNonInteractiveConfirmationReturnsQuestion(t *testing.T) { + streamer := &fakeStreamer{responses: [][]rootly.AiChatStreamEvent{{ + eventOf(rootly.AiChatStreamEventConfirmationRequired, `{"checkpoint_id":"`+testCheckpointID+`","question":"Approve sending the message?","expires_at":"2099-01-01T00:00:00Z"}`), + eventOf(rootly.AiChatStreamEventDone, `{"status":"confirmation_required"}`), + }}} + var stdout, stderr bytes.Buffer + err := RunPrint(context.Background(), streamer, newPrintOptions(&stdout, &stderr)) + if err == nil || !strings.Contains(err.Error(), "Approve sending the message?") { + t.Fatalf("RunPrint() error = %v, want question in error", err) + } +} + +func TestRunPrintAppendsStdinToPrompt(t *testing.T) { + streamer := &fakeStreamer{responses: [][]rootly.AiChatStreamEvent{{ + eventOf(rootly.AiChatStreamEventDone, `{"status":"complete"}`), + }}} + var stdout, stderr bytes.Buffer + options := newPrintOptions(&stdout, &stderr) + options.Prompt = "Explain this" + options.Stdin = strings.NewReader("log contents") + options.IsStdinTTY = func() bool { return false } + + if err := RunPrint(context.Background(), streamer, options); err != nil { + t.Fatalf("RunPrint() error = %v", err) + } + if got, want := *streamer.calls[0].Message, "Explain this\n\nlog contents"; got != want { + t.Errorf("message = %q, want %q", got, want) + } +} + +func newPrintOptions(stdout, stderr *bytes.Buffer) PrintOptions { + return PrintOptions{ + Prompt: "test prompt", + Stdin: strings.NewReader(""), + Stdout: stdout, + Stderr: stderr, + IsStdinTTY: func() bool { return true }, + IsStdoutTTY: func() bool { return false }, + IsStderrTTY: func() bool { return false }, + } +} + +func TestRunPrintPreservesSkillSlashPrompt(t *testing.T) { + prompt := "/incident-retro summarize INC-8" + streamer := &fakeStreamer{responses: [][]rootly.AiChatStreamEvent{{ + eventOf(rootly.AiChatStreamEventDone, `{"status":"complete"}`), + }}} + var stdout, stderr bytes.Buffer + options := newPrintOptions(&stdout, &stderr) + options.Prompt = prompt + if err := RunPrint(context.Background(), streamer, options); err != nil { + t.Fatalf("RunPrint() error = %v", err) + } + if len(streamer.calls) != 1 || streamer.calls[0].Message == nil || *streamer.calls[0].Message != prompt { + t.Errorf("print prompt = %#v, want original prompt %q", streamer.calls, prompt) + } +} diff --git a/internal/agent/sanitize.go b/internal/agent/sanitize.go new file mode 100644 index 0000000..bfc5c0c --- /dev/null +++ b/internal/agent/sanitize.go @@ -0,0 +1,121 @@ +package agent + +import ( + "strings" + "unicode/utf8" +) + +func sanitizeTerminalText(text string) string { + var sanitized strings.Builder + sanitized.Grow(len(text)) + for index := 0; index < len(text); { + character := rune(text[index]) + size := 1 + if character >= 0x80 { + character, size = utf8.DecodeRuneInString(text[index:]) + } + if character == utf8.RuneError && size == 1 && text[index] >= 0x80 && text[index] <= 0x9f { + index += size + continue + } + if character == 0x1b { + index += size + index = skipEscapeSequence(text, index) + continue + } + if character == '\n' || character == '\t' { + sanitized.WriteRune(character) + index += size + continue + } + if character <= 0x1f || character >= 0x7f && character <= 0x9f { + index += size + continue + } + switch character { + case 0x061c, 0x200e, 0x200f: + index += size + continue + } + if character >= 0x202a && character <= 0x202e || character >= 0x2066 && character <= 0x2069 { + index += size + continue + } + sanitized.WriteString(text[index : index+size]) + index += size + } + return sanitized.String() +} + +func skipEscapeSequence(text string, index int) int { + if index >= len(text) { + return index + } + switch text[index] { + case '[': + return skipControlSequence(text, index+1) + case ']', 'P', '^', '_', 'X': + return skipStringControl(text, index+1) + default: + return index + } +} + +func skipControlSequence(text string, index int) int { + for index < len(text) { + character := text[index] + index++ + if character >= 0x40 && character <= 0x7e { + return index + } + } + return index +} + +func skipStringControl(text string, index int) int { + for index < len(text) { + if strings.HasPrefix(text[index:], "\u009c") { + return index + len("\u009c") + } + switch text[index] { + case '\a': + return index + 1 + case 0x1b: + if index+1 < len(text) && text[index+1] == '\\' { + return index + 2 + } + } + index++ + } + return index +} + +type sanitizedTerminalError struct { + message string + cause error +} + +func (err sanitizedTerminalError) Error() string { + return err.message +} + +func (err sanitizedTerminalError) Unwrap() error { + return err.cause +} + +func sanitizeTerminalError(err error) error { + if err == nil { + return nil + } + original := err.Error() + message := sanitizeTerminalText(original) + if message == original { + return err + } + return sanitizedTerminalError{message: message, cause: err} +} + +// SanitizeTerminalError removes terminal controls while preserving the cause. +func SanitizeTerminalError(err error) error { + return sanitizeTerminalError(err) +} diff --git a/internal/agent/sanitize_test.go b/internal/agent/sanitize_test.go new file mode 100644 index 0000000..dbbf0cb --- /dev/null +++ b/internal/agent/sanitize_test.go @@ -0,0 +1,36 @@ +package agent + +import ( + "strings" + "testing" +) + +func TestSanitizeTerminalTextRemovesControlsAndBidiCharacters(t *testing.T) { + input := "\x1b]52;c;ZXZpbA==\x07\x1b[2J\x1b]8;;https://evil.example\x07link\x1b]8;;\x07\r\x7f\u0080\u009f\u061c\u200e\u200f\u202a\u202e\u2066\u2069\n\tcafé 😀" + want := "link\n\tcafé 😀" + + if got := sanitizeTerminalText(input); got != want { + t.Errorf("sanitizeTerminalText() = %q, want %q", got, want) + } +} + +func TestStatusResponseErrorSanitizesServerMessage(t *testing.T) { + err := StatusResponseError(400, []byte(`{"errors":[{"title":"bad \u001b]52;c;ZXZpbA==\u0007 \u001b[2J"}]}`)) + if err == nil { + t.Fatal("StatusResponseError() returned nil") + } + for _, unwanted := range []string{"]52;", "[2J", "\x1b"} { + if strings.Contains(err.Error(), unwanted) { + t.Errorf("error = %q, want no %q", err, unwanted) + } + } +} + +func TestSanitizeTerminalTextRemovesC1ControlsIndividually(t *testing.T) { + if got, want := sanitizeTerminalText("\u009dvisible\u009c"), "visible"; got != want { + t.Errorf("sanitizeTerminalText() = %q, want %q", got, want) + } + if got, want := sanitizeTerminalText(string([]byte{0x9b, 'v', 'i', 's', 'i', 'b', 'l', 'e'})), "visible"; got != want { + t.Errorf("sanitizeTerminalText() with a raw C1 byte = %q, want %q", got, want) + } +} diff --git a/internal/agent/stream.go b/internal/agent/stream.go new file mode 100644 index 0000000..69504b3 --- /dev/null +++ b/internal/agent/stream.go @@ -0,0 +1,200 @@ +package agent + +import ( + "context" + "encoding/json" + "errors" + "fmt" + "iter" + "strings" + "time" + + rootly "github.com/rootlyhq/rootly-go" +) + +type Streamer interface { + StreamAiChatEvents(context.Context, *rootly.StreamAiChatParams) (iter.Seq2[rootly.AiChatStreamEvent, error], error) +} + +type Event struct { + Type string + SessionID *rootly.AiChatStreamSessionID + Text *rootly.AiChatStreamText + Status *rootly.AiChatStreamStatus + TaskUpdate *rootly.AiChatStreamTaskUpdate + SkillLoaded *rootly.AiChatStreamSkillLoaded + ConfirmationRequired *rootly.AiChatStreamConfirmationRequired + CheckpointResolved *rootly.AiChatStreamCheckpointResolved + Error *rootly.AiChatStreamError + Done *rootly.AiChatStreamDone + ClientToolCall *ClientToolCall + ClientToolsAccepted *int +} + +const EventClientToolCall = "client_tool_call" + +// ClientToolCall asks the CLI to run one of the local tools it advertised. +type ClientToolCall struct { + CallID string `json:"call_id"` + Name string `json:"name"` + Arguments map[string]any `json:"arguments"` +} + +const aiChatDisabledMessage = "Rootly AI chat isn't enabled for your organization. Contact your Rootly account team or Rootly support to turn it on." + +func AIChatDisabledError() error { + return aiChatDisabledError{} +} + +func IsAIChatDisabled(err error) bool { + var disabled aiChatDisabledError + return errors.As(err, &disabled) +} + +type aiChatDisabledError struct{} + +func (aiChatDisabledError) Error() string { + return aiChatDisabledMessage +} + +func StreamTurn(ctx context.Context, streamer Streamer, params *rootly.StreamAiChatParams, yield func(Event) bool) error { + events, err := streamer.StreamAiChatEvents(ctx, params) + if err != nil { + return describeError(err) + } + + for event, err := range events { + if err != nil { + return describeError(err) + } + decoded, known, err := decodeEvent(event) + if err != nil { + return err + } + if known && !yield(decoded) { + return nil + } + } + + return nil +} + +func decodeEvent(event rootly.AiChatStreamEvent) (Event, bool, error) { + decoded := Event{Type: event.Type} + var payload any + switch event.Type { + case rootly.AiChatStreamEventSessionID: + payload = &rootly.AiChatStreamSessionID{} + case rootly.AiChatStreamEventText: + payload = &rootly.AiChatStreamText{} + case rootly.AiChatStreamEventStatus: + payload = &rootly.AiChatStreamStatus{} + case rootly.AiChatStreamEventTaskUpdate: + payload = &rootly.AiChatStreamTaskUpdate{} + case rootly.AiChatStreamEventSkillLoaded: + payload = &rootly.AiChatStreamSkillLoaded{} + case rootly.AiChatStreamEventConfirmationRequired: + payload = &rootly.AiChatStreamConfirmationRequired{} + case rootly.AiChatStreamEventCheckpointResolved: + payload = &rootly.AiChatStreamCheckpointResolved{} + case rootly.AiChatStreamEventError: + payload = &rootly.AiChatStreamError{} + case rootly.AiChatStreamEventDone: + payload = &rootly.AiChatStreamDone{} + case EventClientToolCall: + payload = &ClientToolCall{} + default: + return Event{}, false, nil + } + + if err := event.Decode(payload); err != nil { + return Event{}, false, fmt.Errorf("decode %s event: %w", event.Type, err) + } + + switch value := payload.(type) { + case *rootly.AiChatStreamSessionID: + decoded.SessionID = value + var accepted struct { + ClientToolsAccepted *int `json:"client_tools_accepted"` + } + if json.Unmarshal(event.Data, &accepted) == nil { + decoded.ClientToolsAccepted = accepted.ClientToolsAccepted + } + case *rootly.AiChatStreamText: + decoded.Text = value + case *rootly.AiChatStreamStatus: + decoded.Status = value + case *rootly.AiChatStreamTaskUpdate: + decoded.TaskUpdate = value + case *rootly.AiChatStreamSkillLoaded: + decoded.SkillLoaded = value + case *rootly.AiChatStreamConfirmationRequired: + decoded.ConfirmationRequired = value + case *rootly.AiChatStreamCheckpointResolved: + decoded.CheckpointResolved = value + case *rootly.AiChatStreamError: + decoded.Error = value + case *rootly.AiChatStreamDone: + decoded.Done = value + case *ClientToolCall: + if value.CallID == "" || value.Name == "" { + return Event{}, false, fmt.Errorf("decode %s event: missing call_id or name", event.Type) + } + decoded.ClientToolCall = value + } + + return decoded, true, nil +} + +func describeError(err error) error { + var streamErr *rootly.AiChatStreamStatusError + if !errors.As(err, &streamErr) { + return err + } + + status := streamErr.StatusCode + body := strings.TrimSpace(string(streamErr.Body)) + title := "" + if body != "" { + var response struct { + Errors []struct { + Title string `json:"title"` + } `json:"errors"` + } + if json.Unmarshal([]byte(body), &response) == nil && len(response.Errors) > 0 { + title = response.Errors[0].Title + } + } + + switch status { + case 401: + return NotAuthenticatedError("not authenticated: run `rootly login` or set ROOTLY_API_KEY") + case 403: + if strings.Contains(strings.ToLower(title), "not enabled") { + return AIChatDisabledError() + } + if strings.Contains(strings.ToLower(title+" "+body), "ai.chat:write") { + return fmt.Errorf("OAuth token requires ai.chat:write scope for this action. Run `rootly logout && rootly login` to refresh your scopes") + } + case 429: + if retryAfter, ok := streamErr.RetryAfter(); ok { + seconds := int64(retryAfter / time.Second) + if retryAfter%time.Second != 0 { + seconds++ + } + if seconds < 1 { + seconds = 1 + } + return fmt.Errorf("rate limited by Rootly AI: retry in %ds", seconds) + } + return fmt.Errorf("rate limit exceeded (60/min); please try again shortly") + } + + if title != "" { + return fmt.Errorf("%s", title) + } + if body != "" { + return fmt.Errorf("request failed with status %d: %s", status, body) + } + return fmt.Errorf("request failed with status %d", status) +} diff --git a/internal/agent/stream_test.go b/internal/agent/stream_test.go new file mode 100644 index 0000000..bcbd7ee --- /dev/null +++ b/internal/agent/stream_test.go @@ -0,0 +1,155 @@ +package agent + +import ( + "context" + "encoding/json" + "errors" + "fmt" + "iter" + "net/http" + "strings" + "testing" + + rootly "github.com/rootlyhq/rootly-go" +) + +type fakeStreamer struct { + responses [][]rootly.AiChatStreamEvent + calls []*rootly.StreamAiChatParams +} + +func (f *fakeStreamer) StreamAiChatEvents(_ context.Context, params *rootly.StreamAiChatParams) (iter.Seq2[rootly.AiChatStreamEvent, error], error) { + f.calls = append(f.calls, params) + if len(f.responses) == 0 { + return nil, errors.New("no fake stream response") + } + response := f.responses[0] + f.responses = f.responses[1:] + return func(yield func(rootly.AiChatStreamEvent, error) bool) { + for _, event := range response { + if !yield(event, nil) { + return + } + } + }, nil +} + +func eventOf(eventType, data string) rootly.AiChatStreamEvent { + return rootly.AiChatStreamEvent{Type: eventType, Data: json.RawMessage(data)} +} + +func TestStreamTurnDecodesKnownEventsAndIgnoresUnknown(t *testing.T) { + streamer := &fakeStreamer{responses: [][]rootly.AiChatStreamEvent{{ + eventOf(rootly.AiChatStreamEventText, `{"content":"hello"}`), + eventOf("future_event", `{"anything":true}`), + eventOf(rootly.AiChatStreamEventDone, `{"status":"complete"}`), + }}} + var received []Event + if err := StreamTurn(context.Background(), streamer, &rootly.StreamAiChatParams{}, func(event Event) bool { + received = append(received, event) + return true + }); err != nil { + t.Fatalf("StreamTurn() error = %v", err) + } + if len(received) != 2 { + t.Fatalf("received %d events, want 2", len(received)) + } + if received[0].Text == nil || received[0].Text.Content != "hello" { + t.Errorf("text event = %#v, want decoded content", received[0]) + } + if received[1].Done == nil || received[1].Done.Status != "complete" { + t.Errorf("done event = %#v, want decoded status", received[1]) + } +} + +func TestDescribeError(t *testing.T) { + tests := []struct { + name string + err error + want string + }{ + { + name: "unauthorized", + err: fmt.Errorf("wrapped: %w", newStreamStatusError(http.StatusUnauthorized, `{"errors":[{"title":"Invalid token"}]}`, "")), + want: "not authenticated: run `rootly login` or set ROOTLY_API_KEY", + }, + { + name: "oauth ai chat scope", + err: newStreamStatusError(http.StatusForbidden, `{"errors":[{"title":"OAuth token requires ai.chat:write scope for this action"}]}`, ""), + want: "OAuth token requires ai.chat:write scope for this action. Run `rootly logout && rootly login` to refresh your scopes", + }, + { + name: "rate limit with retry after", + err: newStreamStatusError(http.StatusTooManyRequests, "rate limited", "17"), + want: "rate limited by Rootly AI: retry in 17s", + }, + { + name: "rate limit retry after minimum", + err: newStreamStatusError(http.StatusTooManyRequests, "rate limited", "0"), + want: "rate limited by Rootly AI: retry in 1s", + }, + { + name: "rate limit without retry after", + err: newStreamStatusError(http.StatusTooManyRequests, "rate limited", ""), + want: "rate limit exceeded (60/min); please try again shortly", + }, + { + name: "internal server title", + err: newStreamStatusError(http.StatusInternalServerError, `{"errors":[{"title":"Internal server error"}]}`, ""), + want: "Internal server error", + }, + { + name: "internal server body", + err: newStreamStatusError(http.StatusInternalServerError, "temporary failure", ""), + want: "request failed with status 500: temporary failure", + }, + } + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + if got := describeError(tt.err).Error(); got != tt.want { + t.Errorf("describeError() = %q, want %q", got, tt.want) + } + }) + } +} + +func TestDescribeErrorRetryAfterHTTPDate(t *testing.T) { + got := describeError(newStreamStatusError( + http.StatusTooManyRequests, + "rate limited", + "Wed, 21 Oct 2015 07:28:00 GMT", + )).Error() + if got != "rate limited by Rootly AI: retry in 1s" { + t.Errorf("describeError() = %q, want retry in 1s", got) + } +} + +func newStreamStatusError(statusCode int, body, retryAfter string) *rootly.AiChatStreamStatusError { + header := make(http.Header) + if retryAfter != "" { + header.Set("Retry-After", retryAfter) + } + return &rootly.AiChatStreamStatusError{ + StatusCode: statusCode, + Header: header, + Body: []byte(body), + } +} + +func TestStreamTurnDescribesRequestErrors(t *testing.T) { + streamer := streamError{err: newStreamStatusError(http.StatusUnauthorized, "auth failed", "")} + err := StreamTurn(context.Background(), streamer, &rootly.StreamAiChatParams{}, func(Event) bool { + return true + }) + if err == nil || !strings.Contains(err.Error(), "not authenticated") { + t.Fatalf("StreamTurn() error = %v, want authentication guidance", err) + } +} + +type streamError struct { + err error +} + +func (s streamError) StreamAiChatEvents(context.Context, *rootly.StreamAiChatParams) (iter.Seq2[rootly.AiChatStreamEvent, error], error) { + return nil, s.err +} diff --git a/internal/agent/tui.go b/internal/agent/tui.go new file mode 100644 index 0000000..60cdabf --- /dev/null +++ b/internal/agent/tui.go @@ -0,0 +1,1815 @@ +package agent + +import ( + "context" + "errors" + "fmt" + "math/rand/v2" + "net/http" + "strings" + "time" + + "charm.land/bubbles/v2/key" + "charm.land/bubbles/v2/list" + "charm.land/bubbles/v2/spinner" + "charm.land/bubbles/v2/textarea" + "charm.land/bubbles/v2/viewport" + tea "charm.land/bubbletea/v2" + "charm.land/glamour/v2" + "charm.land/lipgloss/v2" + "github.com/charmbracelet/colorprofile" + "github.com/charmbracelet/x/ansi" + "github.com/google/uuid" + "github.com/oapi-codegen/nullable" + openapi_types "github.com/oapi-codegen/runtime/types" + rootly "github.com/rootlyhq/rootly-go" + + "github.com/rootlyhq/rootly-cli/internal/api" + "github.com/rootlyhq/rootly-cli/internal/mcp" + "github.com/rootlyhq/rootly-cli/internal/oauth" +) + +type tuiState int + +const ( + tuiIdle tuiState = iota + tuiStreaming + tuiConfirming +) + +type streamMessage struct { + event Event + err error + done bool + turn int +} + +type TUIOptions struct { + NoColor bool + Version string + API TUIAPI + SessionID *uuid.UUID + SessionTitle string + APIHost string + CurrentUser *rootly.CurrentUserResponse + Incident *api.Incident + Alert *api.Alert + MCP MCPManager + ToolResults ToolResultPoster + WorkingDir string +} + +type TUIAPI interface { + GetCurrentUserWithResponse(context.Context) (*rootly.GetCurrentUserResponse, error) + ListAiChatSessionsWithResponse(context.Context, *rootly.ListAiChatSessionsParams) (*rootly.ListAiChatSessionsResponse, error) + ListAISkillsWithResponse(context.Context, *rootly.ListAISkillsParams, ...rootly.RequestEditorFn) (*rootly.ListAISkillsResponse, error) + ListAiChatSessionMessagesWithResponse(context.Context, openapi_types.UUID, *rootly.ListAiChatSessionMessagesParams) (*rootly.ListAiChatSessionMessagesResponse, error) + UpdateAiChatSessionWithApplicationVndAPIPlusJSONBodyWithResponse(context.Context, rootly.ID, rootly.UpdateAiChatSessionApplicationVndAPIPlusJSONRequestBody) (*rootly.UpdateAiChatSessionResponse, error) + DeleteAiChatSessionWithResponse(context.Context, rootly.ID) (*rootly.DeleteAiChatSessionResponse, error) + GetIncidentByID(context.Context, string) (*api.Incident, error) + GetAlertByID(context.Context, string) (*api.Alert, error) + APIHost() string +} + +type transcriptTurn struct { + user string + assistant string +} + +type transcriptBlock struct { + text string + markdown string + first bool + gapBefore bool + renderedWidth int + rendered string +} + +type historyLoadedMsg struct { + request int + messages []rootly.AiChatSessionMessage + err error + title string + sessionTitle string +} + +type sessionPageMsg struct { + request int + sessions []rootly.AiChatSession + before string + hasMore bool + err error +} + +type sessionActionMsg struct { + request int + action string + message string + sessionTitle *string + err error +} + +type bindingResolvedMsg struct { + request int + incident *api.Incident + alert *api.Alert + err error +} + +type TUIModel struct { + streamer Streamer + state tuiState + composer textarea.Model + spinner spinner.Model + width int + height int + windowSizeReceived bool + viewport viewport.Model + viewportContent string + newOutputAvailable bool + status string + answer strings.Builder + committedOffset int + sessionID *uuid.UUID + pending *rootly.AiChatStreamConfirmationRequired + ctx context.Context + cancel context.CancelFunc + stream chan streamMessage + startedAt time.Time + doneStatus string + doneReceived bool + idleInterrupt bool + currentParams *rootly.StreamAiChatParams + turn int + noColor bool + renderer *glamour.TermRenderer + rendererWidth int + renderedSource string + rendered string + renderCount int + resumeTextSeparator bool + api TUIAPI + apiHost string + currentUser *rootly.CurrentUserResponse + userLoaded bool + sessionTitle string + incident *api.Incident + alert *api.Alert + incidentUsed bool + alertUsed bool + resumedSession bool + historyRequest int + historyPending bool + queuedPrompt queuedSend + promptHistory []string + historyIndex int + historyDraft string + currentPrompt string + turnStartedAt time.Time + turnVerbIndex int + verbIndex func(int) int + now func() time.Time + version string + colorProfile colorprofile.Profile + colorProfileKnown bool + welcomePending bool + welcomeAnimating bool + welcomeFrame int + welcomeGeneration int + quitting bool + lastAnswer string + transcript []transcriptBlock + transcriptCache string + transcriptCacheWidth int + transcriptCacheDirty bool + conversation []transcriptTurn + activities activityTracker + statusRequest int + sessionPicker list.Model + pickerOpen bool + pickerBefore string + pickerHasMore bool + pickerRequest int + pickerLoading bool + deleteConfirming bool + actionRequest int + commandMenu list.Model + commandMenuOpen bool + skillPicker list.Model + skillPickerOpen bool + skillPickerLoading bool + skillPickerRequest int + skillPickerNextPage int + skillPickerQuery string + skillInvocationReq int + skillSuggestionReq int + skillSuggestionCache map[string][]rootly.AiSkill + skillsUnavailable bool + mcp MCPManager + mcpPage mcpPage + toolResults ToolResultPoster + streamCtx context.Context + sentTools int + toolsOffShown bool + localConfirm *localConfirmation + localToolQueue []*localConfirmation + localApprovalRunning bool + localReadOnlyRunning int + workingDir string + terminalBlurred bool + toolsOnShown bool + mentionRequest int + marketRequest int + credHintShown bool + mcpStartupReported bool +} + +type localConfirmation struct { + ref mcp.ToolRef + call *ClientToolCall + approval bool + readOnly bool + options mcp.ApprovalOptions +} + +type localToolDoneMsg struct { + turn int + err error + isError bool + activity *pendingActivity + credHint bool + failure string + tool string + approval bool + readOnly bool +} + +func NewTUIModel(ctx context.Context, streamer Streamer, options TUIOptions) *TUIModel { + composer := textarea.New() + composer.DynamicHeight = true + composer.MinHeight = 1 + composer.MaxHeight = 8 + composer.Placeholder = "Ask Rootly..." + composer.ShowLineNumbers = false + composer.SetHeight(1) + composer.SetPromptFunc(2, func(info textarea.PromptInfo) string { + if info.LineNumber == 0 { + return "❯ " + } + return " " + }) + composer.KeyMap.InsertNewline = key.NewBinding(key.WithKeys("enter", "shift+enter", "alt+enter")) + _ = composer.Focus() + + spin := spinner.New() + spin.Spinner = spinner.Dot + + var currentUser *rootly.CurrentUserResponse + if options.CurrentUser != nil { + user := *options.CurrentUser + currentUser = &user + } + sessionTitle := sanitizeTerminalText(options.SessionTitle) + + model := &TUIModel{ + streamer: streamer, + composer: composer, + spinner: spin, + ctx: ctx, + width: 80, + height: 24, + noColor: options.NoColor, + version: options.Version, + verbIndex: rand.IntN, + now: time.Now, + api: options.API, + apiHost: options.APIHost, + currentUser: currentUser, + userLoaded: options.CurrentUser != nil, + sessionID: options.SessionID, + sessionTitle: normalizeSessionTitle(sessionTitle), + resumedSession: options.SessionID != nil, + incident: options.Incident, + alert: options.Alert, + historyIndex: -1, + skillSuggestionCache: make(map[string][]rootly.AiSkill), + mcp: options.MCP, + toolResults: options.ToolResults, + workingDir: options.WorkingDir, + } + model.refreshMCPViews() + model.viewport = viewport.New() + model.viewport.SoftWrap = true + model.viewport.FillHeight = true + model.viewport.KeyMap = viewport.KeyMap{} + model.resetTranscript() + return model +} + +func (m *TUIModel) Init() tea.Cmd { + m.welcomePending = !m.noColor && !m.colorProfileKnown + m.welcomeAnimating = !m.noColor && m.colorProfileKnown && m.colorProfile >= colorprofile.ANSI256 + m.welcomeFrame = 0 + m.quitting = false + cmds := make([]tea.Cmd, 0, 5) + cmds = append(cmds, m.composer.Focus()) + if m.sessionID != nil { + cmds = append(cmds, m.loadHistory(*m.sessionID)) + } + if m.welcomeAnimating { + cmds = append(cmds, m.welcomeBannerTick()) + } + if watch := m.watchMCP(); watch != nil { + return tea.Batch(tea.Sequence(cmds...), watch, mcpStartupCheck) + } + return tea.Sequence(cmds...) +} + +func (m *TUIModel) Update(message tea.Msg) (tea.Model, tea.Cmd) { + switch msg := message.(type) { + case tea.ColorProfileMsg: + m.colorProfile = msg.Profile + m.colorProfileKnown = true + if !m.welcomePending { + return m, nil + } + m.welcomePending = false + if !m.noColor && msg.Profile >= colorprofile.ANSI256 { + m.welcomeAnimating = true + m.welcomeFrame = 0 + cmd := m.welcomeBannerTick() + return m, cmd + } + return m, nil + case welcomeBannerTickMsg: + if msg.generation != m.welcomeGeneration { + return m, nil + } + if !m.welcomeAnimating { + return m, nil + } + m.welcomeFrame++ + if m.welcomeFrame >= welcomeSettleFrame { + m.welcomeAnimating = false + return m, nil + } + cmd := m.welcomeBannerTick() + return m, cmd + case tea.WindowSizeMsg: + m.handleWindowSize(msg) + return m, nil + case spinner.TickMsg: + if m.state == tuiStreaming { + var cmd tea.Cmd + m.spinner, cmd = m.spinner.Update(msg) + return m, tea.Batch(cmd, m.spinner.Tick) + } + case streamMessage: + return m.updateStream(msg) + case historyLoadedMsg: + cmd := m.onHistoryLoaded(msg) + return m, cmd + case statusLoadedMsg: + if msg.request != m.statusRequest { + return m, nil + } + m.userLoaded = true + if msg.response != nil && msg.response.ApplicationVndAPIJSON200 != nil { + m.currentUser = msg.response.ApplicationVndAPIJSON200 + } + cmd := m.println(m.statusLines()) + return m, cmd + case sessionPageMsg: + cmd := m.handleSessionPage(msg) + return m, cmd + case skillPickerPageMsg, skillInvocationResolvedMsg, skillSuggestionTickMsg, skillSuggestionsLoadedMsg: + return m.updateSkillMessage(msg) + case sessionActionMsg: + if msg.request != m.actionRequest { + return m, nil + } + cmd := m.handleSessionAction(msg) + return m, cmd + case bindingResolvedMsg: + if msg.request != m.actionRequest { + return m, nil + } + if msg.err != nil { + cmd := m.println(sanitizeTerminalText(msg.err.Error())) + return m, cmd + } + m.mentionRequest++ + m.sessionID = nil + m.resetAnswer() + m.resetTranscript() + m.resumedSession = false + m.lastAnswer = "" + m.sessionTitle = "" + m.incident, m.alert = msg.incident, msg.alert + m.incidentUsed, m.alertUsed = false, false + m.status = "" + m.println(faint(m.noColor, m.bindingHeader())) + cmd := m.restartWelcomeAnimation() + return m, cmd + case tea.PasteMsg: + if m.pickerOpen || m.skillPickerOpen { + return m.updatePickerMessage(message) + } + if m.deleteConfirming || m.state == tuiConfirming { + return m, nil + } + if m.idleInterrupt { + m.idleInterrupt = false + m.status = "" + } + m.historyIndex = -1 + var cmd tea.Cmd + m.composer, cmd = m.composer.Update(msg) + menuCmd := m.refreshCommandMenu() + return m, tea.Batch(cmd, menuCmd) + case tea.KeyPressMsg: + return m.updateKey(msg) + default: + return m.updateOtherMessage(message) + } + return m, nil +} + +func (m *TUIModel) updatePickerMessage(message tea.Msg) (tea.Model, tea.Cmd) { + if m.pickerOpen { + var cmd tea.Cmd + m.sessionPicker, cmd = m.sessionPicker.Update(message) + return m, cmd + } + if m.skillPickerOpen { + var cmd tea.Cmd + m.skillPicker, cmd = m.skillPicker.Update(message) + return m, cmd + } + return m, nil +} + +func (m *TUIModel) handleWindowSize(msg tea.WindowSizeMsg) { + m.windowSizeReceived = true + m.width = msg.Width + m.height = msg.Height + m.composer.SetWidth(max(1, msg.Width-2)) + m.composer.MaxHeight = max(3, msg.Height/3) + if m.pickerOpen { + m.sessionPicker.SetSize(max(1, msg.Width), max(1, msg.Height)) + } + if m.skillPickerOpen { + m.skillPicker.SetSize(max(1, msg.Width), max(1, msg.Height)) + } + if m.commandMenuOpen { + m.commandMenu.SetSize(max(20, msg.Width), min(6, len(m.commandMenu.Items()))) + } +} + +func (m *TUIModel) updateSkillMessage(message tea.Msg) (tea.Model, tea.Cmd) { + switch msg := message.(type) { + case skillPickerPageMsg: + cmd := m.handleSkillPickerPage(msg) + return m, cmd + case skillInvocationResolvedMsg: + if msg.request != m.skillInvocationReq { + return m, nil + } + cmd := m.handleSkillInvocation(msg) + return m, cmd + case skillSuggestionTickMsg: + value := strings.TrimSpace(m.composer.Value()) + if msg.request != m.skillSuggestionReq || value != m.composer.Value() || + !strings.HasPrefix(value, "/") || strings.ContainsAny(value, " \t\r\n") || + strings.TrimPrefix(strings.ToLower(value), "/") != msg.prefix { + return m, nil + } + cmd := m.loadSkillSuggestions(msg.prefix, msg.request) + return m, cmd + case skillSuggestionsLoadedMsg: + if msg.request != m.skillSuggestionReq { + return m, nil + } + if msg.statusCode == http.StatusForbidden || msg.statusCode == http.StatusNotFound { + m.skillsUnavailable = true + cmd := m.refreshCommandMenu() + return m, cmd + } + if msg.err != nil { + return m, nil + } + if m.skillSuggestionCache == nil { + m.skillSuggestionCache = make(map[string][]rootly.AiSkill) + } + m.skillSuggestionCache[msg.prefix] = msg.skills + cmd := m.refreshCommandMenu() + return m, cmd + default: + return m, nil + } +} + +func (m *TUIModel) updateStream(message streamMessage) (tea.Model, tea.Cmd) { + if message.turn != m.turn { + return m, nil + } + if m.state != tuiStreaming { + return m, nil + } + if message.done { + return m.finishTurn(message.err) + } + if message.err != nil { + m.status = sanitizeTerminalText(message.err.Error()) + cmd := m.waitForStream() + return m, cmd + } + if message.event.Type == EventClientToolCall { + cmd := m.handleClientToolCall(message.event.ClientToolCall) + return m, tea.Batch(cmd, m.waitForStream()) + } + m.handleEvent(message.event) + m.noteClientToolsAccepted(message.event.ClientToolsAccepted) + if message.event.Type == rootly.AiChatStreamEventError && message.event.Error != nil { + return m.finishTurn(errors.New(sanitizeTerminalText(message.event.Error.Message))) + } + if message.event.Type == rootly.AiChatStreamEventText { + m.commitCompleteAnswerBlocks() + } + cmd := m.waitForStream() + return m, cmd +} + +func (m *TUIModel) handleEvent(event Event) { + for _, line := range m.activities.Handle(event) { + m.println(m.activityLine(line)) + } + switch event.Type { + case rootly.AiChatStreamEventSessionID: + if event.SessionID != nil { + if id, err := uuid.Parse(event.SessionID.SessionID); err == nil { + m.sessionID = &id + } + } + case rootly.AiChatStreamEventText: + if event.Text != nil { + if m.resumeTextSeparator { + if m.answer.Len() > 0 && !strings.HasSuffix(m.answer.String(), "\n") { + m.answer.WriteString("\n\n") + } + m.resumeTextSeparator = false + } + m.answer.WriteString(sanitizeTerminalText(event.Text.Content)) + } + case rootly.AiChatStreamEventStatus: + if event.Status != nil { + m.status = sanitizeTerminalText(event.Status.Status) + } + case rootly.AiChatStreamEventTaskUpdate: + if event.TaskUpdate != nil { + m.status = fmt.Sprintf("%s (%s)", sanitizeTerminalText(event.TaskUpdate.Title), sanitizeTerminalText(event.TaskUpdate.Status)) + } + case rootly.AiChatStreamEventSkillLoaded: + if event.SkillLoaded != nil { + skillName := sanitizeTerminalText(event.SkillLoaded.Skill.Name) + m.println(faint(m.noColor, "• Loaded skill "+skillName)) + } + case rootly.AiChatStreamEventConfirmationRequired: + m.pending = event.ConfirmationRequired + case rootly.AiChatStreamEventError: + if event.Error != nil { + m.status = sanitizeTerminalText(event.Error.Message) + } + case rootly.AiChatStreamEventDone: + m.doneReceived = true + if event.Done != nil { + m.doneStatus = sanitizeTerminalText(event.Done.Status) + } + } +} + +func (m *TUIModel) resetAnswer() { + m.answer.Reset() + m.committedOffset = 0 +} + +func (m *TUIModel) uncommittedAnswer() string { + answer := m.answer.String() + offset := min(max(m.committedOffset, 0), len(answer)) + return answer[offset:] +} + +func (m *TUIModel) commitCompleteAnswerBlocks() { + answer := m.answer.String() + boundary := lastTopLevelBlockBoundary(answer, m.committedOffset) + if boundary <= m.committedOffset { + return + } + chunk := answer[m.committedOffset:boundary] + if strings.TrimSpace(chunk) == "" { + return + } + gapBefore := m.committedOffset > 0 + first := m.committedOffset == 0 + m.committedOffset = boundary + m.appendAssistantTranscriptBlock(chunk, first, gapBefore) +} + +func lastTopLevelBlockBoundary(answer string, committed int) int { + boundary := committed + inFence := false + var fenceMarker rune + fenceLength := 0 + for lineStart := 0; lineStart < len(answer); { + lineEnd := strings.IndexByte(answer[lineStart:], '\n') + if lineEnd < 0 { + break + } + lineEnd += lineStart + line := answer[lineStart:lineEnd] + if marker, length, closing, ok := markdownFence(line); ok { + if !inFence { + inFence = true + fenceMarker, fenceLength = marker, length + } else if closing && marker == fenceMarker && length >= fenceLength { + inFence = false + } + } + if !inFence && lineEnd+1 < len(answer) && answer[lineEnd+1] == '\n' && lineEnd+2 > committed { + boundary = lineEnd + 2 + } + lineStart = lineEnd + 1 + } + return boundary +} + +func markdownFence(line string) (marker rune, length int, closing, ok bool) { + indent := len(line) - len(strings.TrimLeft(line, " ")) + if indent > 3 { + return 0, 0, false, false + } + line = line[indent:] + if line == "" || (line[0] != '`' && line[0] != '~') { + return 0, 0, false, false + } + marker = rune(line[0]) + length = 0 + for length < len(line) && rune(line[length]) == marker { + length++ + } + if length < 3 { + return 0, 0, false, false + } + closing = strings.TrimSpace(line[length:]) == "" + return marker, length, closing, true +} + +func (m *TUIModel) finishTurn(err error) (tea.Model, tea.Cmd) { + if err == nil && !m.doneReceived { + err = errors.New("AI chat stream ended before the turn finished") + } + m.doneReceived = false + if m.cancel != nil { + m.cancel() + } + m.cancel = nil + m.stream = nil + m.state = tuiIdle + m.localConfirm = nil + m.localToolQueue = nil + m.localApprovalRunning = false + m.localReadOnlyRunning = 0 + if err != nil { + for _, line := range m.activities.Finish(false) { + m.println(m.activityLine(line)) + } + } else if m.doneStatus == "confirmation_required" { + m.activities.dropConfirmation() + for _, line := range m.activities.Finish(true) { + m.println(m.activityLine(line)) + } + } else if m.doneStatus != "" { + for _, line := range m.activities.Finish(isSuccessfulDoneStatus(m.doneStatus)) { + m.println(m.activityLine(line)) + } + } + var commands []tea.Cmd + if err != nil { + m.status = "" + m.pending = nil + m.doneStatus = "" + cmd := m.scrollback(sanitizeTerminalText(err.Error()), lipgloss.NewStyle().Foreground(lipgloss.Color("1"))) + m.activities.Reset() + return m, sequenceTUICommands(append(commands, cmd)...) + } + if m.doneStatus == "error" { + m.status = "" + m.pending = nil + m.doneStatus = "" + cmd := m.scrollback("AI chat turn failed", lipgloss.NewStyle().Foreground(lipgloss.Color("1"))) + m.activities.Reset() + return m, sequenceTUICommands(append(commands, cmd)...) + } + if m.doneStatus == "confirmation_required" && m.pending != nil { + if expired(m.pending) { + m.pending = nil + m.status = "" + m.doneStatus = "" + cmd := m.scrollback("Confirmation expired; no action was taken.", lipgloss.NewStyle()) + m.activities.Reset() + return m, sequenceTUICommands(append(commands, cmd)...) + } + m.state = tuiConfirming + if rawAnswer := m.answer.String(); strings.TrimSpace(rawAnswer) != "" { + m.lastAnswer = rawAnswer + if remainder := m.uncommittedAnswer(); strings.TrimSpace(remainder) != "" { + m.appendAssistantTranscriptBlock(remainder, m.committedOffset == 0, m.committedOffset > 0) + } + m.committedOffset = len(rawAnswer) + } + return m, sequenceTUICommands(commands...) + } + if m.doneStatus == "cancelled" { + m.status = "" + m.pending = nil + m.doneStatus = "" + cmd := m.scrollback("Action cancelled.", lipgloss.NewStyle()) + m.activities.Reset() + return m, sequenceTUICommands(append(commands, cmd)...) + } + if m.doneStatus == "expired" { + m.status = "" + m.pending = nil + m.doneStatus = "" + cmd := m.scrollback("Confirmation expired; no action was taken.", lipgloss.NewStyle()) + m.activities.Reset() + return m, sequenceTUICommands(append(commands, cmd)...) + } + + rawAnswer := m.answer.String() + answer := strings.TrimSpace(rawAnswer) + if answer != "" { + m.lastAnswer = rawAnswer + m.conversation = append(m.conversation, transcriptTurn{user: m.currentPrompt, assistant: rawAnswer}) + remainder := m.uncommittedAnswer() + if strings.TrimSpace(remainder) != "" { + m.appendAssistantTranscriptBlock(remainder, m.committedOffset == 0, m.committedOffset > 0) + } + m.printlnWithGap(m.turnFooter(), true) + m.committedOffset = len(rawAnswer) + } + if m.doneStatus != "" && m.doneStatus != "complete" && m.doneStatus != "confirmed" { + m.status = "turn halted: " + sanitizeTerminalText(m.doneStatus) + } else { + m.status = "" + } + m.pending = nil + m.doneStatus = "" + m.turnStartedAt = time.Time{} + m.currentPrompt = "" + m.activities.Reset() + return m, sequenceTUICommands(commands...) +} + +func sequenceTUICommands(commands ...tea.Cmd) tea.Cmd { + sequence := make([]tea.Cmd, 0, len(commands)) + for _, command := range commands { + if command != nil { + sequence = append(sequence, command) + } + } + if len(sequence) == 0 { + return nil + } + return tea.Sequence(sequence...) +} + +func (m *TUIModel) scrollback(message string, style lipgloss.Style) tea.Cmd { + lines := make([]string, 0, 4) + rawAnswer := m.answer.String() + if strings.TrimSpace(rawAnswer) != "" { + m.lastAnswer = rawAnswer + if remainder := m.uncommittedAnswer(); strings.TrimSpace(remainder) != "" { + m.appendAssistantTranscriptBlock(remainder, m.committedOffset == 0, m.committedOffset > 0) + } + } + if m.noColor { + lines = append(lines, message) + } else { + lines = append(lines, style.Render(message)) + } + if !m.turnStartedAt.IsZero() && m.sessionID != nil { + lines = append(lines, "\n"+m.turnFooter()) + } + m.resetAnswer() + m.turnStartedAt = time.Time{} + m.currentPrompt = "" + return m.println(strings.Join(lines, "\n")) +} + +func (m *TUIModel) updateKey(msg tea.KeyPressMsg) (tea.Model, tea.Cmd) { + if m.commandMenuOpen { + if handled, cmd := m.handleCommandMenuKey(msg); handled { + return m, cmd + } + } + if m.mcpPage.open { + return m.handleMCPPageKey(msg) + } + if m.pickerOpen { + return m.handlePickerKey(msg) + } + if m.skillPickerOpen { + return m.handleSkillPickerKey(msg) + } + if m.deleteConfirming { + switch strings.ToLower(msg.String()) { + case "y": + m.deleteConfirming = false + cmd := m.deleteSession() + return m, cmd + case "n", "esc": + m.deleteConfirming = false + return m, nil + } + return m, nil + } + if m.localConfirm != nil { + switch strings.ToLower(msg.String()) { + case "y": + cmd := m.resolveLocalConfirmation(true) + return m, cmd + case "a", "p": + cmd := m.rememberLocalConfirmation(strings.EqualFold(msg.String(), "p")) + return m, cmd + case "n", "esc": + cmd := m.resolveLocalConfirmation(false) + return m, cmd + } + if msg.String() != "ctrl+c" { + return m, nil + } + } + if m.scrollViewportKey(msg) { + return m, nil + } + if msg.String() == "ctrl+c" { + if m.state == tuiStreaming { + m.localConfirm = nil + m.localToolQueue = nil + m.localApprovalRunning = false + m.localReadOnlyRunning = 0 + cmd := m.scrollback("Stopped locally. The server may still finish this turn and save it to the session.", lipgloss.NewStyle().Faint(true)) + if m.cancel != nil { + m.cancel() + } + m.state = tuiIdle + m.cancel = nil + m.stream = nil + m.status = "" + return m, cmd + } + if m.idleInterrupt { + cmd := m.quitCommand() + return m, cmd + } + m.idleInterrupt = true + m.queuedPrompt = queuedSend{} + m.mentionRequest++ + m.composer.Reset() + m.status = "press Ctrl+C again to exit" + return m, nil + } + if m.idleInterrupt { + m.idleInterrupt = false + m.status = "" + } + + if m.state == tuiConfirming { + switch strings.ToLower(msg.String()) { + case "y": + cmd := m.handleConfirmationDecision(rootly.StreamAiChatParamsDecisionConfirm) + return m, cmd + case "n", "esc": + cmd := m.handleConfirmationDecision(rootly.StreamAiChatParamsDecisionCancel) + return m, cmd + } + return m, nil + } + + if msg.String() == "ctrl+d" && strings.TrimSpace(m.composer.Value()) == "" { + cmd := m.quitCommand() + return m, cmd + } + if m.state == tuiStreaming { + return m, nil + } + if m.historyKey(msg) { + return m, nil + } + + if msg.Code == tea.KeyEnter && !msg.Mod.Contains(tea.ModAlt) && !msg.Mod.Contains(tea.ModShift) { + cmd := m.submit() + return m, cmd + } + if msg.Code == tea.KeyEnter && (msg.Mod.Contains(tea.ModAlt) || msg.Mod.Contains(tea.ModShift)) { + m.idleInterrupt = false + var cmd tea.Cmd + m.composer, cmd = m.composer.Update(msg) + menuCmd := m.refreshCommandMenu() + return m, tea.Batch(cmd, menuCmd) + } + + if msg.String() == "esc" { + m.idleInterrupt = false + } + m.historyIndex = -1 + var cmd tea.Cmd + m.composer, cmd = m.composer.Update(msg) + menuCmd := m.refreshCommandMenu() + return m, tea.Batch(cmd, menuCmd) +} + +func (m *TUIModel) handleCommandMenuKey(msg tea.KeyPressMsg) (bool, tea.Cmd) { + switch msg.String() { + case "esc": + m.commandMenuOpen = false + m.skillSuggestionReq++ + return true, nil + case "up", "down", "pgup", "pgdown", "home", "end", "ctrl+home", "ctrl+end": + var cmd tea.Cmd + m.commandMenu, cmd = m.commandMenu.Update(msg) + return true, cmd + case "tab": + if item, ok := m.commandMenu.SelectedItem().(slashCommandItem); ok && item.mention { + m.insertMention(item.command) + } else if ok { + m.composer.SetValue(item.command + " ") + m.composer.CursorEnd() + } + m.commandMenuOpen = false + m.skillSuggestionReq++ + return true, nil + case "enter": + if m.deleteConfirming || m.localConfirm != nil || m.state == tuiConfirming || m.state == tuiStreaming { + return true, nil + } + typed := strings.TrimSpace(m.composer.Value()) + if item, ok := m.commandMenu.SelectedItem().(slashCommandItem); ok && !item.mention && + strings.HasPrefix(strings.ToLower(item.command), strings.ToLower(typed)) { + if item.skill || item.insert { + m.composer.SetValue(item.command + " ") + m.composer.CursorEnd() + m.commandMenuOpen = false + m.skillSuggestionReq++ + cmd := m.composer.Focus() + return true, cmd + } + if len(typed) < len(item.command) { + m.composer.SetValue(item.command) + } + } + m.commandMenuOpen = false + cmd := m.submit() + return true, cmd + default: + return false, nil + } +} + +func (m *TUIModel) handleConfirmationDecision(decision rootly.StreamAiChatParamsDecision) tea.Cmd { + if expired(m.pending) { + m.pending = nil + m.state = tuiIdle + m.status = "" + m.doneStatus = "" + return m.scrollback("Confirmation expired; no action was taken.", lipgloss.NewStyle()) + } + return m.resume(decision) +} + +func (m *TUIModel) submit() tea.Cmd { + if m.continueLine() { + return nil + } + m.settleWelcomeAnimation() + message := strings.TrimSpace(m.composer.Value()) + m.composer.Reset() + m.skillInvocationReq++ + m.skillSuggestionReq++ + m.idleInterrupt = false + m.status = "" + if message == "" { + return nil + } + if isExitAlias(message) { + return m.quitCommand() + } + if strings.HasPrefix(message, "/") { + m.commandMenuOpen = false + return m.runSlashCommand(message) + } + if cmd := m.sendWithResources(message); cmd != nil { + return cmd + } + return m.submitPrompt(message) +} + +func (m *TUIModel) submitPrompt(message string) tea.Cmd { + return m.sendPrompt(message, message, "") +} + +func (m *TUIModel) sendPrompt(echo, message, note string) tea.Cmd { + if m.historyPending { + m.queuedPrompt = queuedSend{echo: echo, message: message, note: note} + m.status = historyLoadingStatus + return nil + } + m.mentionRequest++ + m.currentPrompt = echo + m.promptHistory = append(m.promptHistory, echo) + m.historyIndex = -1 + m.actionRequest++ + m.startNewTurn() + m.activities.Reset() + params := &rootly.StreamAiChatParams{Message: &message, SessionID: m.sessionID} + if m.incident != nil && !m.incidentUsed { + if id, err := uuid.Parse(m.incident.ID); err == nil { + params.IncidentID = &id + } + m.incidentUsed = true + } + if m.alert != nil && !m.alertUsed { + if id, err := uuid.Parse(m.alert.ID); err == nil { + params.AlertID = &id + } + m.alertUsed = true + } + commands := []tea.Cmd{m.println(renderPromptEcho(echo, m.noColor))} + if note != "" { + commands = append(commands, m.println(faint(m.noColor, note))) + } + return sequenceTUICommands(append(commands, m.startStream(params))...) +} + +func (m *TUIModel) scrollViewportKey(msg tea.KeyPressMsg) bool { + switch msg.String() { + case "pgup": + m.viewport.PageUp() + case "pgdown": + m.viewport.PageDown() + case "ctrl+home": + m.viewport.GotoTop() + case "ctrl+end": + m.viewport.GotoBottom() + case "home": + if m.composer.Value() != "" { + return false + } + m.viewport.GotoTop() + case "end": + if m.composer.Value() != "" { + return false + } + m.viewport.GotoBottom() + default: + return false + } + m.updateNewOutputIndicator() + return true +} + +func (m *TUIModel) runSlashCommand(command string) tea.Cmd { + parts := strings.Fields(strings.TrimSpace(command)) + if len(parts) == 0 { + return nil + } + switch parts[0] { + case "/exit", "/quit": + return m.quitCommand() + case "/new", "/clear": + m.startNewSession() + m.status = "Started a new session." + return m.restartWelcomeAnimation() + case "/help": + m.status = m.helpView() + case "/status": + return m.statusCommand() + case "/resume", "/sessions": + return m.openSessionPicker() + case "/mcp", "/connectors": + return m.openMCPPage() + case "/prompts": + return m.listMCPPrompts() + case "/resources": + return m.listMCPResources() + case "/skills": + query := strings.TrimSpace(strings.TrimPrefix(strings.TrimSpace(command), "/skills")) + return m.openSkillPicker(query) + case "/rename": + title := strings.TrimSpace(strings.Join(parts[1:], " ")) + body := rootly.UpdateAiChatSession{} + message := "Renamed to generated title." + if title == "" { + body.Title = nullable.NewNullNullable[string]() + } else { + body.Title = nullable.NewNullableWithValue(title) + message = "Renamed to " + title + "." + } + return m.sessionAction("rename", message, body, &title) + case "/pin", "/unpin": + pinned := parts[0] == "/pin" + body := rootly.UpdateAiChatSession{Pinned: &pinned} + action := "Unpinned" + if pinned { + action = "Pinned" + } + if m.sessionID == nil { + return m.println("No session yet: send a message first.") + } + return m.sessionAction(parts[0], action+" session "+shortSessionID(m.sessionID.String())+".", body, nil) + case "/archive": + archived := true + if m.sessionID == nil { + return m.println("No session yet: send a message first.") + } + return m.sessionAction("archive", "Archived session "+shortSessionID(m.sessionID.String())+".", rootly.UpdateAiChatSession{Archived: &archived}, nil) + case "/delete": + if m.sessionID == nil { + return m.println("No session yet: send a message first.") + } + m.deleteConfirming = true + return nil + case "/copy": + if m.sessionID == nil { + return m.println("No session yet: send a message first.") + } + if m.lastAnswer == "" { + return m.println("Nothing to copy yet.") + } + return tea.Batch(tea.SetClipboard(m.lastAnswer), m.println(faint(m.noColor, "Copied last answer."))) + case "/export": + return m.exportTranscript(strings.Join(parts[1:], " ")) + case "/incident", "/alert": + if len(parts) < 2 { + return m.println(fmt.Sprintf("Usage: %s ", parts[0])) + } + return m.resolveBinding(strings.TrimPrefix(parts[0], "/"), parts[1]) + default: + if cmd, ok := m.runMCPPrompt(command); ok { + return cmd + } + if isSkillInvocationToken(parts[0]) && !isBuiltinSlashCommand(parts[0]) { + return m.lookupSkillForPrompt(parts[0], strings.TrimSpace(command)) + } + return m.println(fmt.Sprintf("Unknown command %s. Type /help.", parts[0])) + } + return nil +} + +func (m *TUIModel) startStream(params *rootly.StreamAiChatParams) tea.Cmd { + m.resetAnswer() + return m.beginStream(params) +} + +func (m *TUIModel) beginStream(params *rootly.StreamAiChatParams) tea.Cmd { + m.turn++ + turn := m.turn + m.state = tuiStreaming + m.resumeTextSeparator = params.CheckpointID != nil + m.pending = nil + m.doneStatus = "" + m.doneReceived = false + if m.turnStartedAt.IsZero() { + m.startNewTurn() + } + m.startedAt = m.turnStartedAt + m.currentParams = params + m.stream = make(chan streamMessage, 16) + ctx, cancel := context.WithCancel(m.ctx) + m.cancel = cancel + stream := m.stream + m.streamCtx = ctx + m.localConfirm = nil + m.localToolQueue = nil + m.localApprovalRunning = false + m.localReadOnlyRunning = 0 + m.sentTools = 0 + streamCtx := ctx + if m.mcp != nil { + tools, _ := m.mcp.ClientTools() + m.sentTools = len(tools) + streamCtx = api.WithClientTools(ctx, tools) + } + + go func() { + defer close(stream) + err := StreamTurn(streamCtx, m.streamer, params, func(event Event) bool { + select { + case stream <- streamMessage{event: event, turn: turn}: + return true + case <-ctx.Done(): + return false + } + }) + select { + case stream <- streamMessage{err: err, done: true, turn: turn}: + case <-ctx.Done(): + } + }() + + return tea.Batch(m.waitForStream(), m.spinner.Tick) +} + +func (m *TUIModel) resume(decision rootly.StreamAiChatParamsDecision) tea.Cmd { + if m.pending == nil || m.sessionID == nil { + m.state = tuiIdle + m.status = "confirmation is missing its session or checkpoint" + return nil + } + checkpointID, err := uuid.Parse(m.pending.CheckpointID) + if err != nil { + m.state = tuiIdle + m.status = "confirmation checkpoint is invalid" + return nil + } + params := &rootly.StreamAiChatParams{ + SessionID: m.sessionID, + CheckpointID: &checkpointID, + Decision: &decision, + } + m.pending = nil + return m.beginStream(params) +} + +func (m *TUIModel) waitForStream() tea.Cmd { + stream := m.stream + turn := m.turn + return func() tea.Msg { + message, ok := <-stream + if !ok { + return streamMessage{err: fmt.Errorf("stream closed"), done: true, turn: turn} + } + return message + } +} + +func (m *TUIModel) View() tea.View { + content := m.viewContent() + view := tea.NewView(content) + view.AltScreen = true + view.ReportFocus = true + return view +} + +func (m *TUIModel) viewContent() string { + if m.quitting { + return "" + } + if m.mcpPage.open { + return m.mcpPageView() + } + if m.pickerOpen { + return m.sessionPicker.View() + } + if m.skillPickerOpen { + return m.skillPicker.View() + } + if !m.windowSizeReceived || m.height <= 0 || m.width <= 0 { + return m.footerContent() + } + + overlay := m.overlayContent() + footerRows := lipgloss.Height(m.footerContent()) + overlayRows := min(lipgloss.Height(overlay), max(0, m.height-footerRows)) + bodyHeight := max(0, m.height-footerRows-overlayRows) + + wasAtBottom := m.viewport.AtBottom() + widthChanged := m.viewport.Width() != m.width + heightChanged := m.viewport.Height() != bodyHeight + if widthChanged { + m.viewport.SetWidth(m.width) + } + if heightChanged { + m.viewport.SetHeight(bodyHeight) + } + body := m.transcriptContent() + if body != m.viewportContent { + m.viewport.SetContent(body) + m.viewportContent = body + if wasAtBottom { + m.viewport.GotoBottom() + m.newOutputAvailable = false + } else { + m.newOutputAvailable = true + } + } else if wasAtBottom && (widthChanged || heightChanged) { + m.viewport.GotoBottom() + } + m.updateNewOutputIndicator() + footer := m.footerContent() + + sections := make([]string, 0, 3) + if bodyHeight > 0 { + sections = append(sections, m.viewport.View()) + } + if overlayRows > 0 { + sections = append(sections, clipViewRows(overlay, overlayRows, max(1, m.width-1))) + } + sections = append(sections, footer) + return strings.Join(sections, "\n") +} + +func (m *TUIModel) transcriptContent() string { + m.rebuildTranscriptCache() + header := strings.TrimSuffix(m.welcomeBannerFrame(), "\n") + if !m.welcomeAnimating { + header = strings.TrimSuffix(m.welcomeBanner(), "\n") + } + live := m.liveSection() + liveGapBefore := m.committedOffset > 0 && strings.TrimSpace(m.liveAnswerRemainder()) != "" + return joinHeaderAndConversation(header, m.transcriptCache, live, liveGapBefore) +} + +func (m *TUIModel) rebuildTranscriptCache() { + width := max(1, m.width-1) + if !m.transcriptCacheDirty && m.transcriptCacheWidth == width { + return + } + blocks := make([]transcriptBlock, 0, max(0, len(m.transcript)-1)) + for index := 1; index < len(m.transcript); index++ { + block := &m.transcript[index] + if block.renderedWidth != width { + if block.markdown != "" { + block.rendered = m.renderAssistantAnswerPart(block.markdown, block.first) + } else { + block.rendered = ansi.Wrap(stripTrailingSpacesPerLine(block.text), width, "") + } + block.renderedWidth = width + } + blocks = append(blocks, transcriptBlock{text: block.rendered, gapBefore: block.gapBefore}) + } + m.transcriptCache = joinTranscriptBlocks(blocks) + m.transcriptCacheWidth = width + m.transcriptCacheDirty = false +} + +func joinTranscriptBlocks(blocks []transcriptBlock) string { + var joined strings.Builder + hasBlock := false + for _, block := range blocks { + text := strings.Trim(block.text, "\n") + if text == "" { + continue + } + if hasBlock { + plain := strings.TrimSpace(ansi.Strip(text)) + if block.gapBefore || strings.HasPrefix(plain, "❯ ") || strings.HasPrefix(plain, "● ") { + joined.WriteString("\n\n") + } else { + joined.WriteByte('\n') + } + } + joined.WriteString(text) + hasBlock = true + } + return joined.String() +} + +func joinHeaderAndConversation(header, committed, live string, liveGapBefore bool) string { + header = strings.TrimRight(header, "\n") + committed = strings.Trim(committed, "\n") + live = strings.Trim(live, "\n") + + var joined strings.Builder + joined.WriteString(header) + if committed != "" { + if header != "" { + joined.WriteString("\n\n") + } + joined.WriteString(committed) + } + if live != "" { + if committed == "" && header != "" { + joined.WriteString("\n\n") + } else if committed != "" || header != "" { + plainLive := strings.TrimSpace(ansi.Strip(live)) + if committed != "" && (liveGapBefore || strings.HasPrefix(plainLive, "● ")) { + joined.WriteString("\n\n") + } else { + joined.WriteByte('\n') + } + } + joined.WriteString(live) + } + return joined.String() +} + +func (m *TUIModel) liveSection() string { + var sections []string + switch m.state { + case tuiStreaming: + markdown := m.liveAnswerRemainder() + if strings.TrimSpace(markdown) != "" { + sections = append(sections, m.renderLiveAssistantAnswerPart(markdown, m.committedOffset == 0)) + } + if m.localConfirm != nil { + sections = append(sections, renderLocalToolCard(m.localConfirm.ref, m.localConfirm.call, m.localConfirm.options, m.width, m.noColor)) + } + case tuiConfirming: + if m.pending != nil { + sections = append(sections, renderConfirmationCard(m.pending, m.width, m.noColor, m.now())) + } + default: + if m.status != "" { + sections = append(sections, m.status) + } + } + return strings.Join(sections, "\n") +} + +func (m *TUIModel) streamingStatusLine() string { + if m.state != tuiStreaming { + return "" + } + elapsed := m.now().Sub(m.startedAt).Truncate(time.Second) + status := m.status + if status == "" { + status = statusVerbs[m.turnVerbIndex].Active + "…" + } + line := fmt.Sprintf("%s %s · %s", m.spinner.View(), status, elapsed) + return ansi.Truncate(line, max(1, m.width-1), "") +} + +func (m *TUIModel) liveAnswerRemainder() string { + markdown := m.uncommittedAnswer() + if !m.noColor { + markdown = trimIncompleteTrailingMarkdownLink(markdown) + } + return markdown +} + +func (m *TUIModel) overlayContent() string { + if m.deleteConfirming { + return "Delete this session? This can't be undone.\ny delete · n/Esc cancel" + } else if m.pickerOpen { + return m.sessionPicker.View() + } else if m.skillPickerOpen { + return m.skillPicker.View() + } else if m.commandMenuOpen { + return m.commandMenu.View() + } + return "" +} + +func (m *TUIModel) footerContent() string { + rule := faint(m.noColor, strings.Repeat("─", max(1, m.width-1))) + hint := "/ for commands · exit to quit" + if m.newOutputAvailable { + hint += " · " + faint(m.noColor, "↓ new output · ctrl+end to jump") + } + hint = ansi.Truncate(hint, max(1, m.width-1), "") + footer := []string{rule, m.composerView(), rule, faint(m.noColor, hint)} + if hidden := m.composer.ScrollYOffset(); hidden > 0 { + footer = append([]string{faint(m.noColor, moreLinesHint(hidden))}, footer...) + } + if statusLine := m.streamingStatusLine(); statusLine != "" { + footer = append([]string{statusLine, ""}, footer...) + } else if mcpLine := m.mcpStatusLine(); mcpLine != "" { + footer = append([]string{ansi.Truncate(mcpLine, max(1, m.width-1), "")}, footer...) + } + return strings.Join(footer, "\n") +} + +func (m *TUIModel) updateNewOutputIndicator() { + if m.viewport.AtBottom() { + m.newOutputAvailable = false + } +} + +func clipViewRows(text string, rows, width int) string { + if rows <= 0 { + return "" + } + lines := strings.Split(ansi.Wrap(text, max(1, width), ""), "\n") + if len(lines) > rows { + lines = lines[:rows] + } + return strings.Join(lines, "\n") +} + +func (m *TUIModel) appendTranscriptBlock(text string) { + m.appendTranscriptBlockWithGap(text, false) +} + +func (m *TUIModel) appendTranscriptBlockWithGap(text string, gapBefore bool) { + m.transcript = append(m.transcript, transcriptBlock{text: text, gapBefore: gapBefore}) + m.transcriptCacheDirty = true +} + +func (m *TUIModel) appendAssistantTranscriptBlock(markdown string, first, gapBefore bool) { + if strings.TrimSpace(markdown) == "" { + return + } + m.transcript = append(m.transcript, transcriptBlock{ + markdown: markdown, + first: first, + gapBefore: gapBefore, + }) + m.transcriptCacheDirty = true +} + +func stripTrailingSpacesPerLine(text string) string { + lines := strings.Split(text, "\n") + for index, line := range lines { + lines[index] = strings.TrimRight(line, " ") + } + return strings.Join(lines, "\n") +} + +func (m *TUIModel) resetTranscript() { + m.transcript = []transcriptBlock{{text: ""}} + m.transcriptCache = "" + m.transcriptCacheWidth = 0 + m.transcriptCacheDirty = true + m.conversation = nil +} + +func (m *TUIModel) println(text string) tea.Cmd { + m.appendTranscriptBlock(text) + return nil +} + +func (m *TUIModel) printlnWithGap(text string, gapBefore bool) { + m.appendTranscriptBlockWithGap(text, gapBefore) +} + +func (m *TUIModel) Transcript() string { + m.rebuildTranscriptCache() + header := strings.TrimSuffix(m.welcomeBanner(), "\n") + live := "" + liveGapBefore := false + if remainder := m.uncommittedAnswer(); strings.TrimSpace(remainder) != "" { + live = m.renderLiveAssistantAnswerPart(remainder, m.committedOffset == 0) + liveGapBefore = m.committedOffset > 0 + } + return joinHeaderAndConversation(header, m.transcriptCache, live, liveGapBefore) + "\n" +} + +func (m *TUIModel) renderMarkdown(markdown string, width int, streaming bool) string { + markdown = sanitizeTerminalText(markdown) + if strings.TrimSpace(markdown) == "" { + return "" + } + if width < 20 { + width = 20 + } + apiHost := m.apiHost + if apiHost == "" && m.api != nil { + apiHost = m.api.APIHost() + } + sessionID := "" + if m.sessionID != nil { + sessionID = m.sessionID.String() + } + webBase := "" + if apiHost != "" { + webBase = oauth.DeriveAuthBaseURL(apiHost) + } + opts := markdownOpts{ + width: width + 2, + indent: 2, + noColor: m.noColor, + streaming: streaming, + webBase: webBase, + sessionID: sessionID, + } + prepared, charts := prepareAssistantMarkdown(markdown, opts) + if m.noColor { + return expandAssistantMarkdownCharts(prepared, charts) + } + if m.renderer == nil || m.rendererWidth != width { + renderer, err := glamour.NewTermRenderer(glamour.WithEnvironmentConfig(), glamour.WithWordWrap(width)) + if err != nil { + return expandAssistantMarkdownCharts(prepared, charts) + } + m.renderer = renderer + m.rendererWidth = width + m.renderedSource = "" + } + renderKey := strings.Join([]string{markdown, fmt.Sprint(width), sessionID, apiHost, fmt.Sprint(streaming)}, "\x00") + if m.renderedSource != renderKey { + source, links := preprocessMarkdownLinksWithURLs(prepared, apiHost) + rendered, err := m.renderer.Render(source) + if err != nil { + return expandAssistantMarkdownCharts(prepared, charts) + } + m.rendered = expandAssistantMarkdownCharts(strings.TrimSpace(restoreMarkdownLinks(rendered, links)), charts) + m.renderedSource = renderKey + m.renderCount++ + } + return m.rendered +} + +func (m *TUIModel) handleClientToolCall(call *ClientToolCall) tea.Cmd { + if call == nil { + return nil + } + var tools LocalTools + if m.mcp != nil { + tools = m.mcp + } + ref, known := resolveLocalCall(tools, call) + if !known { + return m.runClientTool(call, false) + } + m.println(faint(m.noColor, localToolActivity(ref))) + confirmation := &localConfirmation{ + ref: ref, + call: call, + approval: tools.NeedsApproval(call.Name, call.Arguments), + } + confirmation.readOnly = !confirmation.approval && localCallReadOnly(tools, ref, call) + if confirmation.approval { + if memory, ok := tools.(LocalApprovalMemory); ok { + confirmation.options = memory.ApprovalOptions(call.Name, call.Arguments) + } + if m.localConfirm != nil || m.localApprovalRunning || m.localReadOnlyRunning > 0 || len(m.localToolQueue) > 0 { + m.localToolQueue = append(m.localToolQueue, confirmation) + return nil + } + m.localConfirm = confirmation + return nil + } + if !confirmation.readOnly { + if m.localConfirm != nil || m.localApprovalRunning || m.localReadOnlyRunning > 0 || len(m.localToolQueue) > 0 { + m.localToolQueue = append(m.localToolQueue, confirmation) + return nil + } + m.localApprovalRunning = true + return m.runApprovedClientTool(call, false) + } + if m.localConfirm != nil || m.localApprovalRunning || len(m.localToolQueue) > 0 { + m.localToolQueue = append(m.localToolQueue, confirmation) + return nil + } + m.localReadOnlyRunning++ + return m.runReadOnlyClientTool(call) +} + +// localCallReadOnly reports whether a call that needs no approval only reads +// on its own merits. A write allowed by a remembered approval or rule is not +// read-only: it still runs one at a time. +func localCallReadOnly(tools LocalTools, ref mcp.ToolRef, call *ClientToolCall) bool { + if reader, ok := tools.(LocalReadOnlyCalls); ok { + return reader.ReadOnlyCall(call.Name, call.Arguments) + } + return ref.ReadOnly +} + +func (m *TUIModel) resolveLocalConfirmation(approved bool) tea.Cmd { + pending := m.localConfirm + m.localConfirm = nil + if pending == nil { + return nil + } + if !approved { + m.println(faint(m.noColor, "Declined local tool "+sanitizeTerminalText(pending.ref.ServerName)+" · "+sanitizeTerminalText(pending.ref.Tool))) + } + m.localApprovalRunning = true + return m.runApprovedClientTool(pending.call, !approved) +} + +// rememberLocalConfirmation approves the pending call and remembers it for +// the session (a) or as a saved allow rule (p), when the card offers that. +func (m *TUIModel) rememberLocalConfirmation(persist bool) tea.Cmd { + pending := m.localConfirm + if pending == nil || (persist && !pending.options.Persist) || (!persist && !pending.options.Session) { + return nil + } + memory, ok := m.mcp.(LocalApprovalMemory) + if !ok { + return nil + } + label := "this tool" + if len(pending.options.Prefix) > 0 { + label = "`" + sanitizeTerminalText(strings.Join(pending.options.Prefix, " ")) + "`" + } + switch err := memory.RememberApproval(pending.call.Name, pending.call.Arguments, persist); { + case err != nil: + m.println(faint(m.noColor, "Could not remember approval: "+sanitizeTerminalError(err).Error())) + case persist: + m.println(faint(m.noColor, "Always allowing "+label+" (saved to ~/.rootly-cli/rules/default.rules)")) + default: + m.println(faint(m.noColor, "Allowing "+label+" for this session")) + } + return m.resolveLocalConfirmation(true) +} + +func (m *TUIModel) runClientTool(call *ClientToolCall, declined bool) tea.Cmd { + return m.runLocalTool(call, declined, false, false) +} + +func (m *TUIModel) runReadOnlyClientTool(call *ClientToolCall) tea.Cmd { + return m.runLocalTool(call, false, false, true) +} + +func (m *TUIModel) runApprovedClientTool(call *ClientToolCall, declined bool) tea.Cmd { + return m.runLocalTool(call, declined, true, false) +} + +func (m *TUIModel) runLocalTool(call *ClientToolCall, declined, approval, readOnly bool) tea.Cmd { + var tools LocalTools + if m.mcp != nil { + tools = m.mcp + } + poster := m.toolResults + ctx := m.streamCtx + if ctx == nil { + ctx = m.ctx + } + turn := m.turn + sessionID := "" + if m.sessionID != nil { + sessionID = m.sessionID.String() + } + activity := m.activities.pending + return func() tea.Msg { + result := mcp.CallResult{Output: userDeclinedOutput, IsError: true} + if !declined { + result = runLocalCall(ctx, tools, call) + } + done := localToolDoneMsg{ + turn: turn, + err: postLocalResult(ctx, poster, sessionID, call, result), + isError: result.IsError, + activity: activity, + credHint: looksLikeExpiredCreds(result), + approval: approval, + readOnly: readOnly, + } + if done.err != nil { + done.tool = localToolLabel(tools, call) + } + if !declined && result.IsError { + done.failure = localToolFailure(tools, call, result.Output) + } + return done + } +} + +func (m *TUIModel) onHistoryLoaded(msg historyLoadedMsg) tea.Cmd { + if msg.request != m.historyRequest { + return nil + } + m.historyPending = false + if m.status == historyLoadingStatus { + m.status = "" + } + if title := normalizeSessionTitle(msg.sessionTitle); title != "" { + m.sessionTitle = title + } + queued := m.queuedPrompt + m.queuedPrompt = queuedSend{} + m.handleHistoryLoaded(msg) + if queued == (queuedSend{}) { + return nil + } + return m.sendPrompt(queued.echo, queued.message, queued.note) +} + +type queuedSend struct { + echo, message, note string +} + +func (m *TUIModel) startNextLocalTool() tea.Cmd { + if m.localConfirm != nil || m.localApprovalRunning || m.localReadOnlyRunning > 0 || len(m.localToolQueue) == 0 { + return nil + } + for _, queued := range m.localToolQueue { + if queued.approval && m.mcp != nil && !m.mcp.NeedsApproval(queued.call.Name, queued.call.Arguments) { + queued.approval = false + } + } + var commands []tea.Cmd + for len(m.localToolQueue) > 0 && !m.localToolQueue[0].approval && m.localToolQueue[0].readOnly { + next := m.localToolQueue[0] + m.localToolQueue = m.localToolQueue[1:] + m.localReadOnlyRunning++ + commands = append(commands, m.runReadOnlyClientTool(next.call)) + } + if len(commands) > 0 { + if len(commands) == 1 { + return commands[0] + } + return tea.Batch(commands...) + } + next := m.localToolQueue[0] + m.localToolQueue = m.localToolQueue[1:] + if !next.approval { + m.localApprovalRunning = true + return m.runApprovedClientTool(next.call, false) + } + m.localConfirm = next + return nil +} diff --git a/internal/agent/tui_helpers.go b/internal/agent/tui_helpers.go new file mode 100644 index 0000000..17e8361 --- /dev/null +++ b/internal/agent/tui_helpers.go @@ -0,0 +1,1283 @@ +package agent + +import ( + "context" + "errors" + "fmt" + "io" + "net/http" + "os" + "path/filepath" + "regexp" + "sort" + "strconv" + "strings" + "time" + + "charm.land/bubbles/v2/help" + "charm.land/bubbles/v2/key" + "charm.land/bubbles/v2/list" + tea "charm.land/bubbletea/v2" + "charm.land/lipgloss/v2" + "github.com/charmbracelet/x/ansi" + "github.com/google/uuid" + rootly "github.com/rootlyhq/rootly-go" + + "github.com/rootlyhq/rootly-cli/internal/api" +) + +type statusLoadedMsg struct { + request int + response *rootly.GetCurrentUserResponse +} + +type skillPickerPageMsg struct { + request int + page int + skills []rootly.AiSkill + nextPage int + statusCode int + err error +} + +type skillInvocationResolvedMsg struct { + request int + token string + prompt string + skill rootly.AiSkill + found bool + statusCode int + err error +} + +type skillSuggestionTickMsg struct { + request int + prefix string +} + +type skillSuggestionsLoadedMsg struct { + request int + prefix string + skills []rootly.AiSkill + statusCode int + err error +} + +type sessionPickerItem struct { + session rootly.AiChatSession + loadMore bool +} + +type skillPickerItem struct { + skill rootly.AiSkill + loadMore bool +} + +type slashCommandItem struct { + command string + description string + skill bool + insert bool + mention bool +} + +func (i slashCommandItem) FilterValue() string { return i.command } +func (i slashCommandItem) Title() string { return i.command } +func (i slashCommandItem) Description() string { return i.description } + +type slashCommandDelegate struct{} + +func (slashCommandDelegate) Render(writer io.Writer, model list.Model, index int, item list.Item) { + command, ok := item.(slashCommandItem) + if !ok { + return + } + prefix := " " + style := lipgloss.NewStyle().Foreground(lipgloss.Color("7")) + if index == model.Index() { + prefix = "› " + style = lipgloss.NewStyle().Foreground(lipgloss.Color("6")) + } + _, _ = fmt.Fprintf(writer, "%s%s %s", prefix, command.command, style.Render(command.description)) +} + +func (slashCommandDelegate) Height() int { return 1 } +func (slashCommandDelegate) Spacing() int { return 0 } +func (slashCommandDelegate) Update(tea.Msg, *list.Model) tea.Cmd { return nil } + +func (i sessionPickerItem) FilterValue() string { + if i.loadMore { + return "Load more" + } + title := sanitizeTerminalText(i.session.Title) + if strings.TrimSpace(title) == "" { + return "Untitled" + } + return title +} + +func (i sessionPickerItem) Title() string { + if i.loadMore { + return "Load more…" + } + title := sanitizeTerminalText(i.session.Title) + if strings.TrimSpace(title) == "" { + return "Untitled" + } + return title +} + +func (i sessionPickerItem) Description() string { + if i.loadMore { + return "Fetch the next page" + } + description := "last active " + relativeAge(i.session.LastActiveAt) + if i.session.Pinned != nil && *i.session.Pinned { + description += " · pinned" + } + return description +} + +func (i skillPickerItem) FilterValue() string { + if i.loadMore { + return "Load more" + } + return strings.Join([]string{ + sanitizeTerminalText(skillString(i.skill.Slug)), + sanitizeTerminalText(skillString(i.skill.Name)), + sanitizeTerminalText(strings.ReplaceAll(i.skill.Description.GetOrEmpty(), "\n", " ")), + }, " ") +} + +func (i skillPickerItem) Title() string { + if i.loadMore { + return "Load more" + } + marker := "○" + if skillEnabled(i.skill) { + marker = "✔" + } + return sanitizeTerminalText(marker + " /" + skillString(i.skill.Slug)) +} + +func (i skillPickerItem) Description() string { + if i.loadMore { + return "Fetch the next page" + } + parts := make([]string, 0, 3) + if name := sanitizeTerminalText(strings.ReplaceAll(skillString(i.skill.Name), "\n", " ")); name != "" { + parts = append(parts, name) + } + if version, err := i.skill.CurrentVersion.Get(); err == nil { + parts = append(parts, fmt.Sprintf("v%d", version)) + } + if invocable := skillInvocableLabels(i.skill.InvocableBy); invocable != "" { + parts = append(parts, invocable) + } + description := strings.Join(parts, " · ") + if text := strings.ReplaceAll(i.skill.Description.GetOrEmpty(), "\n", " "); text != "" { + if description != "" { + description += " · " + } + description += text + } + return sanitizeTerminalText(description) +} + +var skillInvocationTokenPattern = regexp.MustCompile(`^/[a-z0-9]+(?:-[a-z0-9]+)*$`) + +var builtinSkillSlugs = map[string]struct{}{ + "exit": {}, "quit": {}, "new": {}, "clear": {}, "resume": {}, "sessions": {}, + "rename": {}, "pin": {}, "unpin": {}, "archive": {}, "delete": {}, "copy": {}, + "export": {}, "incident": {}, "alert": {}, "status": {}, "help": {}, "skills": {}, + "mcp": {}, "connectors": {}, "prompts": {}, "resources": {}, +} + +func isSkillInvocationToken(token string) bool { + return skillInvocationTokenPattern.MatchString(token) +} + +func isBuiltinSlashCommand(command string) bool { + _, ok := builtinSkillSlugs[strings.TrimPrefix(command, "/")] + return ok +} + +func skillEnabled(skill rootly.AiSkill) bool { + return skill.Enabled != nil && *skill.Enabled +} + +func skillString(value *string) string { + if value == nil { + return "" + } + return *value +} + +func skillInvocableByUser(skill rootly.AiSkill) bool { + for _, invocableBy := range skill.InvocableBy { + if invocableBy == rootly.AiSkillInvocableBy("user") { + return true + } + } + return false +} + +func skillInvocableLabels(values []rootly.AiSkillInvocableBy) string { + labels := make([]string, 0, len(values)) + for _, value := range values { + switch value { + case rootly.AiSkillInvocableBy("user"): + labels = append(labels, "you") + case rootly.AiSkillInvocableBy("model"): + labels = append(labels, "model") + case rootly.AiSkillInvocableBy("automation"): + labels = append(labels, "automations") + } + } + return strings.Join(labels, " + ") +} + +func relativeAge(value time.Time) string { + elapsed := time.Since(value) + switch { + case elapsed < time.Minute: + return "just now" + case elapsed < time.Hour: + return fmt.Sprintf("%dm ago", int(elapsed.Minutes())) + case elapsed < 24*time.Hour: + return fmt.Sprintf("%dh ago", int(elapsed.Hours())) + case elapsed < 30*24*time.Hour: + return fmt.Sprintf("%dd ago", int(elapsed.Hours()/24)) + default: + return value.Format("2006-01-02") + } +} + +func (m *TUIModel) historyKey(msg tea.KeyPressMsg) bool { + if msg.Code != tea.KeyUp && msg.Code != tea.KeyDown { + return false + } + value := m.composer.Value() + browsing := m.historyIndex >= 0 + if msg.Code == tea.KeyUp { + if value != "" && (!browsing || m.composer.Line() != 0) { + return false + } + if !browsing { + m.historyDraft = value + m.historyIndex = len(m.promptHistory) + } + if m.historyIndex > 0 { + m.historyIndex-- + m.composer.SetValue(m.promptHistory[m.historyIndex]) + } + return true + } + if !browsing || (value != "" && m.composer.Line() != m.composer.LineCount()-1) { + return false + } + if m.historyIndex < len(m.promptHistory)-1 { + m.historyIndex++ + m.composer.SetValue(m.promptHistory[m.historyIndex]) + } else { + m.historyIndex = len(m.promptHistory) + m.composer.SetValue(m.historyDraft) + } + return true +} + +func (m *TUIModel) quitCommand() tea.Cmd { + m.settleWelcomeAnimation() + m.quitting = true + if m.sessionID != nil { + m.println(faint(m.noColor, "Resume with: rootly agent --resume "+m.sessionID.String())) + } + return tea.Quit +} + +func (m *TUIModel) settleWelcomeAnimation() { + m.welcomePending = false + m.welcomeAnimating = false +} + +func isExitAlias(message string) bool { + message = strings.TrimSpace(message) + return !strings.ContainsAny(message, "\r\n") && (strings.EqualFold(message, "exit") || strings.EqualFold(message, "quit")) +} + +func renderPromptEcho(prompt string, noColor bool) string { + lines := strings.Split(prompt, "\n") + for index, line := range lines { + prefix := " " + if index == 0 { + prefix = "❯ " + if !noColor { + prefix = lipgloss.NewStyle().Foreground(lipgloss.Color("8")).Render(prefix) + } + } + lines[index] = prefix + line + } + return strings.Join(lines, "\n") +} + +func (m *TUIModel) renderAssistantAnswer(markdown string) string { + return m.renderAssistantAnswerPart(markdown, true) +} + +func (m *TUIModel) renderAssistantAnswerPart(markdown string, first bool) string { + return m.renderAssistantAnswerPartWithStreaming(markdown, first, false) +} + +func (m *TUIModel) renderLiveAssistantAnswerPart(markdown string, first bool) string { + return m.renderAssistantAnswerPartWithStreaming(markdown, first, m.state == tuiStreaming) +} + +func (m *TUIModel) renderAssistantAnswerPartWithStreaming(markdown string, first, streaming bool) string { + lines := strings.Split(m.renderMarkdown(markdown, max(20, m.width-3), streaming), "\n") + start, end := 0, len(lines) + for start < end && strings.TrimSpace(ansi.Strip(lines[start])) == "" { + start++ + } + for end > start && strings.TrimSpace(ansi.Strip(lines[end-1])) == "" { + end-- + } + lines = lines[start:end] + if len(lines) == 0 { + return "" + } + for index, line := range lines { + line = trimLeadingANSIWhitespace(line) + if index == 0 && first { + lines[index] = "● " + line + } else { + lines[index] = " " + line + } + } + return ansi.Wrap(strings.Join(lines, "\n"), max(1, m.width-1), " ") +} + +func trimLeadingANSIWhitespace(line string) string { + var output strings.Builder + started := false + for index := 0; index < len(line); { + if line[index] == '\x1b' && index+1 < len(line) { + end := index + 2 + switch line[index+1] { + case '[': + for end < len(line) && (line[end] < '@' || line[end] > '~') { + end++ + } + if end < len(line) { + end++ + } + case ']': + for end < len(line) && line[end] != '\a' { + end++ + } + if end < len(line) { + end++ + } + } + output.WriteString(line[index:end]) + index = end + continue + } + if !started && (line[index] == ' ' || line[index] == '\t') { + index++ + continue + } + started = true + output.WriteByte(line[index]) + index++ + } + return output.String() +} + +func (m *TUIModel) turnFooter() string { + now := m.now() + duration := now.Sub(m.turnStartedAt).Round(time.Second) + verb := statusVerbs[m.turnVerbIndex].Done + return faint(m.noColor, fmt.Sprintf("✻ %s for %s · done %s", verb, formatTurnDuration(duration), now.Format("3:04 PM"))) +} + +func (m *TUIModel) startNewTurn() { + m.turnStartedAt = m.now() + m.turnVerbIndex = m.verbIndex(len(statusVerbs)) +} + +func formatTurnDuration(duration time.Duration) string { + if duration < 0 { + duration = 0 + } + if duration < time.Minute { + return fmt.Sprintf("%ds", int(duration/time.Second)) + } + return fmt.Sprintf("%dm %ds", int(duration/time.Minute), int(duration/time.Second)%60) +} + +func shortSessionID(id string) string { + if len(id) > 8 { + return id[:8] + } + return id +} + +func faint(noColor bool, text string) string { + if noColor { + return text + } + return lipgloss.NewStyle().Faint(true).Render(text) +} + +func (m *TUIModel) bindingHeader() string { + if m.incident != nil { + id := m.incident.SequentialID + if id == "" { + id = shortSessionID(m.incident.ID) + } + return sanitizeTerminalText(fmt.Sprintf("Chatting about %s: %s", id, m.incident.Title)) + } + if m.alert != nil { + id := m.alert.ShortID + if id == "" { + id = shortSessionID(m.alert.ID) + } + return sanitizeTerminalText(fmt.Sprintf("Chatting about %s: %s", id, m.alert.Summary)) + } + return "" +} + +func (m *TUIModel) helpView() string { + model := help.New() + groups := [][]key.Binding{{ + key.NewBinding(key.WithKeys("enter"), key.WithHelp("enter", "send")), + key.NewBinding(key.WithKeys("shift+enter", "alt+enter"), key.WithHelp("shift/alt+enter", "new line")), + key.NewBinding(key.WithKeys("up", "down"), key.WithHelp("↑/↓", "history")), + }} + return model.FullHelpView(groups) + "\n" + + "/exit, /quit, exit, quit: exit\n" + + "/new, /clear: start a new session\n" + + "/resume, /sessions: choose a session\n" + + "/skills [query]: browse skills\n" + + "/mcp, /connectors: manage MCP servers\n" + + "/prompts, /resources: list MCP prompts and resources (@server:uri attaches one)\n" + + "/rename [title]: rename or reset title\n" + + "/pin, /unpin: change pinned state\n" + + "/archive: archive this session\n" + + "/delete: delete this session\n" + + "/copy: copy last answer\n" + + "/export [file]: write transcript\n" + + "/incident , /alert : bind a resource\n" + + "/status: show session and API info\n" + + "PgUp/PgDn scroll · Ctrl+Home/Ctrl+End scroll · Home/End scroll when composer is empty\n" + + "/help: show this help" +} + +func (m *TUIModel) statusCommand() tea.Cmd { + if m.userLoaded || m.api == nil { + return m.println(m.statusLines()) + } + m.statusRequest++ + request := m.statusRequest + apiClient := m.api + ctx := m.ctx + return func() tea.Msg { + response, _ := apiClient.GetCurrentUserWithResponse(ctx) + return statusLoadedMsg{request: request, response: response} + } +} + +func (m *TUIModel) statusLines() string { + session := "new session" + if m.sessionID != nil { + session = m.sessionID.String() + } + user := "unknown" + if m.currentUser != nil { + attributes := m.currentUser.Data.Attributes + name := currentUserName(attributes) + user = sanitizeTerminalText(strings.TrimSpace(name + " (" + attributes.Email + ")")) + } + host := m.apiHost + if host == "" && m.api != nil { + host = m.api.APIHost() + } + lines := []string{"session " + session} + if title := sanitizeTerminalText(m.sessionTitle); title != "" { + lines = append(lines, "title "+title) + } + lines = append(lines, "user "+user, "API "+sanitizeTerminalText(host)) + if header := m.bindingHeader(); header != "" { + lines = append(lines, "binding "+strings.TrimPrefix(header, "Chatting about ")) + } + for index := range lines { + lines[index] = faint(m.noColor, lines[index]) + } + return strings.Join(lines, "\n") +} + +func (m *TUIModel) loadHistory(sessionID uuid.UUID) tea.Cmd { + m.historyRequest++ + m.historyPending = true + m.queuedPrompt = queuedSend{} + request := m.historyRequest + apiClient := m.api + ctx := m.ctx + return func() tea.Msg { + if apiClient == nil { + return historyLoadedMsg{request: request, err: errors.New("history API unavailable")} + } + messages, sessionTitle, err := loadSessionHistory(ctx, apiClient, sessionID) + return historyLoadedMsg{request: request, messages: messages, err: err, sessionTitle: sessionTitle} + } +} + +func (m *TUIModel) printHistory(messages []rootly.AiChatSessionMessage) { + for _, message := range messages { + content := sanitizeTerminalText(message.Content) + switch message.Role { + case rootly.AiChatSessionMessageRole("user"): + m.println(renderPromptEcho(content, m.noColor)) + m.conversation = append(m.conversation, transcriptTurn{user: content}) + case rootly.AiChatSessionMessageRole("assistant"): + m.appendAssistantTranscriptBlock(content, true, false) + if len(m.conversation) > 0 && m.conversation[len(m.conversation)-1].assistant == "" { + m.conversation[len(m.conversation)-1].assistant = content + } else { + m.conversation = append(m.conversation, transcriptTurn{assistant: content}) + } + } + } +} + +func (m *TUIModel) openSessionPicker() tea.Cmd { + m.pickerOpen = true + m.pickerLoading = true + m.actionRequest++ + m.pickerBefore = "" + m.pickerHasMore = false + delegate := list.NewDefaultDelegate() + m.sessionPicker = list.New(nil, delegate, max(1, m.width), max(1, m.height)) + m.sessionPicker.Title = "Sessions" + m.sessionPicker.SetShowTitle(true) + m.sessionPicker.SetShowHelp(true) + m.sessionPicker.SetShowStatusBar(true) + m.sessionPicker.SetShowPagination(true) + m.sessionPicker.SetFilteringEnabled(true) + m.sessionPicker.SetSize(max(1, m.width), max(1, m.height)) + return m.loadSessionPage("") +} + +func (m *TUIModel) loadSessionPage(before string) tea.Cmd { + if m.api == nil { + m.pickerOpen = false + return m.println("session API unavailable") + } + m.pickerRequest++ + request := m.pickerRequest + apiClient := m.api + ctx := m.ctx + return func() tea.Msg { + limit, archived := 20, false + params := &rootly.ListAiChatSessionsParams{Limit: &limit, Archived: &archived} + if before != "" { + params.Before = &before + } + response, err := apiClient.ListAiChatSessionsWithResponse(ctx, params) + if err != nil { + return sessionPageMsg{request: request, err: err} + } + if response == nil { + return sessionPageMsg{request: request, err: errors.New("empty session list response")} + } + if response.ApplicationVndAPIJSON200 == nil { + return sessionPageMsg{request: request, err: StatusResponseError(response.StatusCode(), response.Body)} + } + data := response.ApplicationVndAPIJSON200 + return sessionPageMsg{ + request: request, sessions: data.Sessions, hasMore: data.Meta.Pagination.HasMore, + before: data.Meta.Pagination.NextBefore.GetOrEmpty(), + } + } +} + +func (m *TUIModel) handleSessionPage(message sessionPageMsg) tea.Cmd { + if message.request != m.pickerRequest || !m.pickerOpen { + return nil + } + m.pickerLoading = false + if message.err != nil { + m.pickerOpen = false + return m.println(sanitizeTerminalText(message.err.Error())) + } + m.pickerBefore = message.before + m.pickerHasMore = message.hasMore && message.before != "" + items := make([]list.Item, 0, len(message.sessions)+1) + for _, session := range message.sessions { + items = append(items, sessionPickerItem{session: session}) + } + if m.pickerHasMore { + items = append(items, sessionPickerItem{loadMore: true}) + } + if len(m.sessionPicker.Items()) == 0 { + m.sessionPicker.SetItems(items) + } else { + for index := len(m.sessionPicker.Items()) - 1; index >= 0; index-- { + if item, ok := m.sessionPicker.Items()[index].(sessionPickerItem); ok && item.loadMore { + m.sessionPicker.RemoveItem(index) + break + } + } + for _, item := range items { + m.sessionPicker.InsertItem(len(m.sessionPicker.Items()), item) + } + } + return nil +} + +func (m *TUIModel) selectPickerItem(item sessionPickerItem) tea.Cmd { + if item.loadMore { + m.pickerLoading = true + return m.loadSessionPage(m.pickerBefore) + } + id, err := uuid.Parse(item.session.ID.String()) + if err != nil { + m.pickerOpen = false + return m.println("invalid session UUID") + } + m.pickerOpen = false + m.pickerLoading = false + m.actionRequest++ + m.mentionRequest++ + m.sessionID = &id + m.sessionTitle = normalizeSessionTitle(item.Title()) + m.resumedSession = true + m.resetAnswer() + m.lastAnswer = "" + m.resetTranscript() + m.incident, m.alert = nil, nil + m.incidentUsed, m.alertUsed = false, false + return tea.Batch(m.loadHistory(id), m.restartWelcomeAnimation()) +} + +func (m *TUIModel) handlePickerKey(msg tea.KeyPressMsg) (tea.Model, tea.Cmd) { + if msg.String() == "esc" { + if m.sessionPicker.FilterState() != list.Unfiltered { + var cmd tea.Cmd + m.sessionPicker, cmd = m.sessionPicker.Update(msg) + return m, cmd + } + m.pickerOpen = false + m.pickerLoading = false + m.pickerRequest++ + return m, nil + } + if msg.String() == "enter" { + selected, ok := m.sessionPicker.SelectedItem().(sessionPickerItem) + if !ok { + return m, nil + } + cmd := m.selectPickerItem(selected) + return m, cmd + } + var cmd tea.Cmd + m.sessionPicker, cmd = m.sessionPicker.Update(msg) + return m, cmd +} + +const aiSkillsDisabledMessage = "AI Skills isn't enabled for this organization." + +func (m *TUIModel) openSkillPicker(query string) tea.Cmd { + m.skillPickerOpen = true + m.skillPickerLoading = true + m.skillPickerNextPage = 0 + m.skillPickerQuery = query + delegate := list.NewDefaultDelegate() + m.skillPicker = list.New(nil, delegate, max(1, m.width), max(1, m.height)) + m.skillPicker.Title = "AI Skills" + m.skillPicker.SetShowTitle(true) + m.skillPicker.SetShowHelp(true) + m.skillPicker.SetShowStatusBar(true) + m.skillPicker.SetShowPagination(true) + m.skillPicker.SetFilteringEnabled(true) + m.skillPicker.SetSize(max(1, m.width), max(1, m.height)) + return m.loadSkillPickerPage(1) +} + +func (m *TUIModel) loadSkillPickerPage(page int) tea.Cmd { + if m.api == nil { + m.skillPickerOpen = false + m.skillPickerLoading = false + return m.println("AI Skills API unavailable.") + } + m.skillPickerRequest++ + request, apiClient, ctx, query := m.skillPickerRequest, m.api, m.ctx, m.skillPickerQuery + sortBy := rootly.ListAISkillsParamsSortName + params := &rootly.ListAISkillsParams{Sort: &sortBy} + if query != "" { + params.FilterSearch = &query + } + return func() tea.Msg { + response, err := apiClient.ListAISkillsWithResponse(ctx, params, aiSkillsPageRequestEditor(page, 25)) + if err != nil { + return skillPickerPageMsg{request: request, page: page, err: err} + } + if response == nil { + return skillPickerPageMsg{request: request, page: page, err: errors.New("empty AI skills response")} + } + if response.ApplicationVndAPIJSON200 == nil { + return skillPickerPageMsg{ + request: request, page: page, statusCode: response.StatusCode(), + err: aiSkillsResponseError(response.StatusCode(), response.Body), + } + } + data := response.ApplicationVndAPIJSON200 + skills := make([]rootly.AiSkill, 0, len(data.Data)) + for _, item := range data.Data { + skills = append(skills, item.Attributes) + } + return skillPickerPageMsg{ + request: request, page: page, skills: skills, + nextPage: aiSkillsNextPage(data.Meta), + } + } +} + +func (m *TUIModel) handleSkillPickerPage(message skillPickerPageMsg) tea.Cmd { + if message.request != m.skillPickerRequest || !m.skillPickerOpen { + return nil + } + m.skillPickerLoading = false + if message.err != nil { + m.skillPickerOpen = false + if message.statusCode == http.StatusForbidden || message.statusCode == http.StatusNotFound { + m.skillsUnavailable = true + m.skillPickerRequest++ + return m.println(aiSkillsDisabledMessage) + } + return m.println(sanitizeTerminalText(message.err.Error())) + } + + m.skillPickerNextPage = message.nextPage + items := make([]list.Item, 0, len(message.skills)+1) + for _, skill := range message.skills { + slug := skillString(skill.Slug) + if slug == "" || !isSkillInvocationToken("/"+slug) || isBuiltinSlashCommand("/"+slug) { + continue + } + items = append(items, skillPickerItem{skill: skill}) + } + if m.skillPickerNextPage > 0 { + items = append(items, skillPickerItem{loadMore: true}) + } + if message.page == 1 { + m.skillPicker.SetItems(items) + return nil + } + for index := len(m.skillPicker.Items()) - 1; index >= 0; index-- { + if item, ok := m.skillPicker.Items()[index].(skillPickerItem); ok && item.loadMore { + m.skillPicker.RemoveItem(index) + break + } + } + for _, item := range items { + m.skillPicker.InsertItem(len(m.skillPicker.Items()), item) + } + return nil +} + +func (m *TUIModel) selectSkillPickerItem(item skillPickerItem) tea.Cmd { + if item.loadMore { + m.skillPickerLoading = true + return m.loadSkillPickerPage(m.skillPickerNextPage) + } + m.skillPickerOpen = false + m.skillPickerLoading = false + m.skillPickerRequest++ + slug := skillString(item.skill.Slug) + token := "/" + slug + if !skillEnabled(item.skill) { + return m.println(token + " is disabled.") + } + if !skillInvocableByUser(item.skill) { + return m.println(token + " can only be loaded by the model, not by name.") + } + draft := strings.TrimLeft(m.composer.Value(), " \t\r\n") + value := token + " " + if draft != "" { + value += draft + } + m.composer.SetValue(value) + m.composer.CursorEnd() + return m.composer.Focus() +} + +func (m *TUIModel) handleSkillPickerKey(msg tea.KeyPressMsg) (tea.Model, tea.Cmd) { + if msg.String() == "esc" { + if m.skillPicker.FilterState() != list.Unfiltered { + var cmd tea.Cmd + m.skillPicker, cmd = m.skillPicker.Update(msg) + return m, cmd + } + m.skillPickerOpen = false + m.skillPickerLoading = false + m.skillPickerRequest++ + return m, nil + } + if msg.String() == "enter" { + selected, ok := m.skillPicker.SelectedItem().(skillPickerItem) + if !ok { + return m, nil + } + cmd := m.selectSkillPickerItem(selected) + return m, cmd + } + var cmd tea.Cmd + m.skillPicker, cmd = m.skillPicker.Update(msg) + return m, cmd +} + +func (m *TUIModel) lookupSkillForPrompt(token, prompt string) tea.Cmd { + if m.api == nil { + return m.println(fmt.Sprintf("Unknown command %s. Type /help.", token)) + } + m.skillInvocationReq++ + request, apiClient, ctx := m.skillInvocationReq, m.api, m.ctx + slug := strings.TrimPrefix(token, "/") + sortBy := rootly.ListAISkillsParamsSortName + params := &rootly.ListAISkillsParams{FilterSlug: &slug, Sort: &sortBy} + return func() tea.Msg { + response, err := apiClient.ListAISkillsWithResponse(ctx, params, aiSkillsPageRequestEditor(1, 25)) + if err != nil { + return skillInvocationResolvedMsg{request: request, token: token, prompt: prompt, err: err} + } + if response == nil { + return skillInvocationResolvedMsg{request: request, token: token, prompt: prompt, err: errors.New("empty AI skills response")} + } + if response.ApplicationVndAPIJSON200 == nil { + return skillInvocationResolvedMsg{ + request: request, token: token, prompt: prompt, statusCode: response.StatusCode(), + err: aiSkillsResponseError(response.StatusCode(), response.Body), + } + } + for _, item := range response.ApplicationVndAPIJSON200.Data { + if skillString(item.Attributes.Slug) == slug { + return skillInvocationResolvedMsg{ + request: request, token: token, prompt: prompt, skill: item.Attributes, found: true, + } + } + } + return skillInvocationResolvedMsg{request: request, token: token, prompt: prompt} + } +} + +func (m *TUIModel) handleSkillInvocation(message skillInvocationResolvedMsg) tea.Cmd { + restorePrompt := func() { + if m.composer.Value() == "" { + m.composer.SetValue(message.prompt) + m.composer.CursorEnd() + } + } + if message.statusCode == http.StatusForbidden || message.statusCode == http.StatusNotFound { + m.skillsUnavailable = true + restorePrompt() + return m.println(fmt.Sprintf("Unknown command %s. Type /help.", message.token)) + } + if message.err != nil { + restorePrompt() + return m.println(fmt.Sprintf("Couldn't check skill %s: %s", message.token, sanitizeTerminalText(message.err.Error()))) + } + if !message.found { + restorePrompt() + return m.println(fmt.Sprintf("Unknown command %s. Type /help.", message.token)) + } + if !skillEnabled(message.skill) { + restorePrompt() + return m.println(message.token + " is disabled.") + } + if !skillInvocableByUser(message.skill) { + restorePrompt() + return m.println(message.token + " can only be loaded by the model, not by name.") + } + return m.submitPrompt(message.prompt) +} + +func aiSkillsPageRequestEditor(page, size int) rootly.RequestEditorFn { + return func(_ context.Context, request *http.Request) error { + if request.URL == nil { + return errors.New("AI skills request has no URL") + } + query := request.URL.Query() + query.Set("page[number]", strconv.Itoa(page)) + query.Set("page[size]", strconv.Itoa(size)) + request.URL.RawQuery = query.Encode() + return nil + } +} + +func aiSkillsNextPage(meta map[string]interface{}) int { + switch value := meta["next_page"].(type) { + case float64: + return int(value) + case int: + return value + case int64: + return int(value) + case string: + page, _ := strconv.Atoi(value) + return page + default: + return 0 + } +} + +func aiSkillsResponseError(status int, body []byte) error { + if status == http.StatusForbidden || status == http.StatusNotFound { + return fmt.Errorf("%s", aiSkillsDisabledMessage) + } + return StatusResponseError(status, body) +} + +func (m *TUIModel) loadSkillSuggestions(prefix string, request int) tea.Cmd { + if m.api == nil || m.skillsUnavailable { + return nil + } + apiClient, ctx := m.api, m.ctx + enabled := true + invocableBy := "user" + sortBy := rootly.ListAISkillsParamsSortName + params := &rootly.ListAISkillsParams{ + FilterEnabled: &enabled, FilterInvocableBy: &invocableBy, Sort: &sortBy, + } + params.FilterSearch = &prefix + return func() tea.Msg { + response, err := apiClient.ListAISkillsWithResponse(ctx, params, aiSkillsPageRequestEditor(1, 6)) + if err != nil { + return skillSuggestionsLoadedMsg{request: request, prefix: prefix, err: err} + } + if response == nil { + return skillSuggestionsLoadedMsg{request: request, prefix: prefix, err: errors.New("empty AI skills response")} + } + if response.ApplicationVndAPIJSON200 == nil { + return skillSuggestionsLoadedMsg{ + request: request, prefix: prefix, statusCode: response.StatusCode(), + err: aiSkillsResponseError(response.StatusCode(), response.Body), + } + } + skills := make([]rootly.AiSkill, 0, len(response.ApplicationVndAPIJSON200.Data)) + for _, item := range response.ApplicationVndAPIJSON200.Data { + skills = append(skills, item.Attributes) + } + return skillSuggestionsLoadedMsg{request: request, prefix: prefix, skills: skills} + } +} + +func (m *TUIModel) refreshCommandMenu() tea.Cmd { + m.skillSuggestionReq++ + request := m.skillSuggestionReq + rawValue := m.composer.Value() + value := strings.TrimSpace(rawValue) + if m.refreshMentionMenu(rawValue) { + return nil + } + if rawValue != value || value == "" || !strings.HasPrefix(value, "/") || strings.ContainsAny(value, " \t\r\n") { + m.commandMenuOpen = false + return nil + } + commands := []slashCommandItem{ + {command: "/exit", description: "exit chat"}, {command: "/quit", description: "exit chat"}, + {command: "/new", description: "new session"}, {command: "/clear", description: "new session"}, + {command: "/resume", description: "choose a session"}, {command: "/sessions", description: "browse sessions"}, + {command: "/skills", description: "browse skills"}, {command: "/mcp", description: "manage MCP servers"}, + {command: "/connectors", description: "manage MCP servers"}, {command: "/prompts", description: "list MCP prompts"}, + {command: "/resources", description: "list MCP resources"}, {command: "/rename", description: "rename session"}, + {command: "/pin", description: "pin session"}, {command: "/unpin", description: "unpin session"}, + {command: "/archive", description: "archive session"}, {command: "/delete", description: "delete session"}, + {command: "/copy", description: "copy last answer"}, {command: "/export", description: "export transcript"}, + {command: "/incident", description: "bind incident"}, {command: "/alert", description: "bind alert"}, + {command: "/status", description: "show session info"}, {command: "/help", description: "show help"}, + } + items := make([]list.Item, 0, 6) + for _, command := range commands { + if strings.HasPrefix(strings.ToLower(command.command), strings.ToLower(value)) { + items = append(items, command) + if len(items) == 6 { + break + } + } + } + prefix := strings.TrimPrefix(strings.ToLower(value), "/") + if skills, cached := m.skillSuggestionCache[prefix]; cached { + for _, skill := range skills { + slug := skillString(skill.Slug) + if slug == "" || !skillEnabled(skill) || !skillInvocableByUser(skill) || + isBuiltinSlashCommand("/"+slug) || !strings.HasPrefix("/"+strings.ToLower(slug), strings.ToLower(value)) { + continue + } + items = append(items, slashCommandItem{ + command: "/" + slug, description: sanitizeTerminalText(strings.ReplaceAll(skillString(skill.Name), "\n", " ")), skill: true, + }) + } + } + items = m.appendPromptItems(items, value) + if len(items) == 0 { + m.commandMenuOpen = false + } else { + m.showCommandMenu(items) + } + if m.api == nil || m.skillsUnavailable { + return nil + } + if _, cached := m.skillSuggestionCache[prefix]; cached { + return nil + } + return tea.Tick(150*time.Millisecond, func(time.Time) tea.Msg { + return skillSuggestionTickMsg{request: request, prefix: prefix} + }) +} + +func (m *TUIModel) sessionAction(action, message string, body rootly.UpdateAiChatSessionApplicationVndAPIPlusJSONRequestBody, sessionTitle *string) tea.Cmd { + if m.sessionID == nil { + return m.println("No session yet: send a message first.") + } + if m.api == nil { + return m.println("session API unavailable") + } + m.actionRequest++ + request, id, apiClient, ctx := m.actionRequest, *m.sessionID, m.api, m.ctx + return func() tea.Msg { + response, err := apiClient.UpdateAiChatSessionWithApplicationVndAPIPlusJSONBodyWithResponse(ctx, rootly.ID(id.String()), body) + if err == nil && response == nil { + err = errors.New("empty session update response") + } else if err == nil && (response.StatusCode() < 200 || response.StatusCode() >= 300) { + err = sessionActionStatusError(response.StatusCode(), response.Body) + } + return sessionActionMsg{request: request, action: action, message: message, sessionTitle: sessionTitle, err: err} + } +} + +func (m *TUIModel) deleteSession() tea.Cmd { + if m.sessionID == nil { + return m.println("No session yet: send a message first.") + } + if m.api == nil { + return m.println("session API unavailable") + } + m.actionRequest++ + request, id, apiClient, ctx := m.actionRequest, *m.sessionID, m.api, m.ctx + return func() tea.Msg { + response, err := apiClient.DeleteAiChatSessionWithResponse(ctx, rootly.ID(id.String())) + if err == nil && response == nil { + err = errors.New("empty session delete response") + } else if err == nil && (response.StatusCode() < 200 || response.StatusCode() >= 300) { + err = sessionActionStatusError(response.StatusCode(), response.Body) + } + return sessionActionMsg{request: request, action: "delete", message: "Deleted session " + shortSessionID(id.String()) + ".", err: err} + } +} + +var errSessionActionNotFound = errors.New("session action not found") + +const sessionActionNotFoundMessage = "Session not found. Incident and alert chats can't be renamed, pinned, archived or deleted from the CLI yet." + +func sessionActionStatusError(status int, body []byte) error { + if status == http.StatusNotFound { + return errSessionActionNotFound + } + return StatusResponseError(status, body) +} + +func (m *TUIModel) handleSessionAction(message sessionActionMsg) tea.Cmd { + if message.err != nil { + if errors.Is(message.err, errSessionActionNotFound) { + return m.println(sessionActionNotFoundMessage) + } + return m.println(sanitizeTerminalText(message.err.Error())) + } + if message.action == "rename" { + m.sessionTitle = "" + if message.sessionTitle != nil { + m.sessionTitle = sanitizeTerminalText(*message.sessionTitle) + } + } + text := sanitizeTerminalText(message.message) + if message.action == "archive" || message.action == "delete" { + m.startNewSession() + m.println(faint(m.noColor, text)) + return m.restartWelcomeAnimation() + } + return m.println(faint(m.noColor, text)) +} + +func (m *TUIModel) startNewSession() { + m.actionRequest++ + m.pickerRequest++ + m.historyRequest++ + m.mentionRequest++ + m.historyPending = false + m.queuedPrompt = queuedSend{} + m.sessionID = nil + m.sessionTitle = "" + m.resetAnswer() + m.resetTranscript() + m.resumedSession = false + m.lastAnswer = "" + m.currentPrompt = "" + m.incident = nil + m.alert = nil + m.incidentUsed = false + m.alertUsed = false +} + +func (m *TUIModel) exportTranscript(filename string) tea.Cmd { + if m.sessionID == nil { + return m.println("No session yet: send a message first.") + } + if len(m.conversation) == 0 { + return m.println("Nothing to export yet.") + } + if strings.TrimSpace(filename) == "" { + filename = fmt.Sprintf("rootly-chat-%s-%s.md", shortSessionID(m.sessionID.String()), time.Now().Format("20060102-150405")) + } + path, err := filepath.Abs(filename) + if err == nil { + var file *os.File + file, err = os.OpenFile(path, os.O_WRONLY|os.O_CREATE|os.O_EXCL, 0600) + if err == nil { + _, err = file.WriteString(m.transcriptMarkdown()) + closeErr := file.Close() + if err == nil { + err = closeErr + } + } + } + if err != nil { + return m.println("Export failed: " + sanitizeTerminalText(err.Error())) + } + return m.println(faint(m.noColor, "Exported "+path)) +} + +func (m *TUIModel) transcriptMarkdown() string { + var output strings.Builder + for _, turn := range m.conversation { + if turn.user != "" { + fmt.Fprintf(&output, "## You\n\n%s\n\n", turn.user) + } + if turn.assistant != "" { + fmt.Fprintf(&output, "## Rootly\n\n%s\n\n", turn.assistant) + } + } + return strings.TrimRight(output.String(), "\n") + "\n" +} + +func (m *TUIModel) resolveBinding(kind, value string) tea.Cmd { + if m.api == nil { + return m.println("resource API unavailable") + } + m.actionRequest++ + request, apiClient, ctx := m.actionRequest, m.api, m.ctx + return func() tea.Msg { + if kind == "incident" { + incident, err := apiClient.GetIncidentByID(ctx, api.NormalizeIncidentID(value)) + return bindingResolvedMsg{request: request, incident: incident, err: err} + } + alert, err := apiClient.GetAlertByID(ctx, value) + return bindingResolvedMsg{request: request, alert: alert, err: err} + } +} + +func normalizeSessionTitle(title string) string { + title = strings.TrimSpace(sanitizeTerminalText(title)) + if title == "" || title == "Untitled" { + return "" + } + return title +} + +func loadSessionHistory(ctx context.Context, apiClient TUIAPI, sessionID uuid.UUID) ([]rootly.AiChatSessionMessage, string, error) { + page, size := 1, 100 + first, err := apiClient.ListAiChatSessionMessagesWithResponse(ctx, sessionID, &rootly.ListAiChatSessionMessagesParams{PageNumber: &page, PageSize: &size}) + if err != nil { + return nil, "", err + } + if first == nil { + return nil, "", errors.New("empty session history response") + } + if first.StatusCode() == 404 { + return nil, "", errSessionNotFound + } + if first.ApplicationVndAPIJSON200 == nil { + return nil, "", StatusResponseError(first.StatusCode(), first.Body) + } + sessionTitle := normalizeSessionTitle(first.ApplicationVndAPIJSON200.Title) + messages := append([]rootly.AiChatSessionMessage(nil), first.ApplicationVndAPIJSON200.Messages...) + pages := historyPageCount(first.ApplicationVndAPIJSON200.Meta) + if pages > 1 { + last, lastSize := pages, 100 + response, err := apiClient.ListAiChatSessionMessagesWithResponse(ctx, sessionID, &rootly.ListAiChatSessionMessagesParams{PageNumber: &last, PageSize: &lastSize}) + if err != nil { + return nil, "", err + } + if response == nil { + return nil, "", errors.New("empty session history response") + } + if response.ApplicationVndAPIJSON200 == nil { + return nil, "", StatusResponseError(response.StatusCode(), response.Body) + } + seen := make(map[string]bool, len(messages)) + for _, message := range messages { + seen[message.ID.String()] = true + } + for _, message := range response.ApplicationVndAPIJSON200.Messages { + if !seen[message.ID.String()] { + messages = append(messages, message) + } + } + } + sort.Slice(messages, func(i, j int) bool { return messages[i].CreatedAt.Before(messages[j].CreatedAt) }) + filtered := messages[:0] + for _, message := range messages { + if message.Role == rootly.AiChatSessionMessageRole("user") || message.Role == rootly.AiChatSessionMessageRole("assistant") { + filtered = append(filtered, message) + } + } + if len(filtered) > 6 { + filtered = filtered[len(filtered)-6:] + } + return filtered, sessionTitle, nil +} + +func historyPageCount(meta map[string]interface{}) int { + pagination, ok := meta["pagination"].(map[string]interface{}) + if !ok { + return 1 + } + for _, key := range []string{"pages", "last"} { + switch value := pagination[key].(type) { + case float64: + if int(value) > 1 { + return int(value) + } + case int: + if value > 1 { + return value + } + } + } + return 1 +} + +func StatusResponseError(status int, body []byte) error { + return sanitizeTerminalError(describeError(&rootly.AiChatStreamStatusError{StatusCode: status, Body: body})) +} + +func (m *TUIModel) showCommandMenu(items []list.Item) { + m.commandMenu = list.New(items, slashCommandDelegate{}, max(20, m.width), min(8, len(items))) + m.commandMenu.SetShowTitle(false) + m.commandMenu.SetShowFilter(false) + m.commandMenu.SetShowStatusBar(false) + m.commandMenu.SetShowPagination(false) + m.commandMenu.SetShowHelp(false) + m.commandMenu.SetFilteringEnabled(false) + m.commandMenuOpen = true +} diff --git a/internal/agent/tui_test.go b/internal/agent/tui_test.go new file mode 100644 index 0000000..0ff9798 --- /dev/null +++ b/internal/agent/tui_test.go @@ -0,0 +1,2629 @@ +package agent + +import ( + "bytes" + "context" + "errors" + "fmt" + "iter" + "net/http" + "net/url" + "os" + "path/filepath" + "reflect" + "strings" + "testing" + "time" + + "charm.land/bubbles/v2/list" + tea "charm.land/bubbletea/v2" + "charm.land/lipgloss/v2" + "github.com/charmbracelet/colorprofile" + "github.com/charmbracelet/x/ansi" + "github.com/google/uuid" + "github.com/oapi-codegen/nullable" + openapi_types "github.com/oapi-codegen/runtime/types" + rootly "github.com/rootlyhq/rootly-go" + + "github.com/rootlyhq/rootly-cli/internal/api" +) + +type blockingStreamer struct { + started chan struct{} +} + +type fakeTUIAPI struct { + sessionPages []*rootly.ListAiChatSessionsResponse + listParams []*rootly.ListAiChatSessionsParams + skillPages []*rootly.ListAISkillsResponse + skillParams []*rootly.ListAISkillsParams + skillQueries []url.Values + skillErrors []error + messagePages []*rootly.ListAiChatSessionMessagesResponse + messageParams []*rootly.ListAiChatSessionMessagesParams + updateBodies []rootly.UpdateAiChatSessionApplicationVndAPIPlusJSONRequestBody + updateStatus int + deleteCalls int + deleteStatus int +} + +func (f *fakeTUIAPI) GetCurrentUserWithResponse(context.Context) (*rootly.GetCurrentUserResponse, error) { + return nil, nil +} + +func (f *fakeTUIAPI) ListAiChatSessionsWithResponse(_ context.Context, params *rootly.ListAiChatSessionsParams) (*rootly.ListAiChatSessionsResponse, error) { + f.listParams = append(f.listParams, params) + if len(f.sessionPages) == 0 { + return testSessionPage(nil, false, ""), nil + } + response := f.sessionPages[0] + f.sessionPages = f.sessionPages[1:] + return response, nil +} + +func (f *fakeTUIAPI) ListAISkillsWithResponse(_ context.Context, params *rootly.ListAISkillsParams, editors ...rootly.RequestEditorFn) (*rootly.ListAISkillsResponse, error) { + f.skillParams = append(f.skillParams, params) + request := &http.Request{URL: &url.URL{}} + for _, editor := range editors { + if err := editor(context.Background(), request); err != nil { + return nil, err + } + } + query := request.URL.Query() + if params.FilterSearch != nil { + query.Set("filter[search]", *params.FilterSearch) + } + if params.FilterSlug != nil { + query.Set("filter[slug]", *params.FilterSlug) + } + if params.FilterEnabled != nil { + query.Set("filter[enabled]", fmt.Sprint(*params.FilterEnabled)) + } + if params.FilterInvocableBy != nil { + query.Set("filter[invocable_by]", *params.FilterInvocableBy) + } + if params.Sort != nil { + query.Set("sort", string(*params.Sort)) + } + request.URL.RawQuery = query.Encode() + f.skillQueries = append(f.skillQueries, request.URL.Query()) + if len(f.skillErrors) > 0 { + err := f.skillErrors[0] + f.skillErrors = f.skillErrors[1:] + return nil, err + } + if len(f.skillPages) == 0 { + return testSkillPage(nil, nil), nil + } + response := f.skillPages[0] + f.skillPages = f.skillPages[1:] + return response, nil +} + +func (f *fakeTUIAPI) ListAiChatSessionMessagesWithResponse(_ context.Context, _ openapi_types.UUID, params *rootly.ListAiChatSessionMessagesParams) (*rootly.ListAiChatSessionMessagesResponse, error) { + f.messageParams = append(f.messageParams, params) + if len(f.messagePages) == 0 { + return nil, errors.New("history not configured") + } + response := f.messagePages[0] + f.messagePages = f.messagePages[1:] + return response, nil +} + +func (f *fakeTUIAPI) UpdateAiChatSessionWithApplicationVndAPIPlusJSONBodyWithResponse(_ context.Context, _ rootly.ID, body rootly.UpdateAiChatSessionApplicationVndAPIPlusJSONRequestBody) (*rootly.UpdateAiChatSessionResponse, error) { + f.updateBodies = append(f.updateBodies, body) + status := f.updateStatus + if status == 0 { + status = http.StatusOK + } + return &rootly.UpdateAiChatSessionResponse{HTTPResponse: &http.Response{StatusCode: status}, Body: []byte(`{"errors":[{"title":"ai.chat:write required"}]}`)}, nil +} + +func (f *fakeTUIAPI) DeleteAiChatSessionWithResponse(context.Context, rootly.ID) (*rootly.DeleteAiChatSessionResponse, error) { + f.deleteCalls++ + status := f.deleteStatus + if status == 0 { + status = http.StatusNoContent + } + return &rootly.DeleteAiChatSessionResponse{ + HTTPResponse: &http.Response{StatusCode: status}, + Body: []byte(`{"errors":[{"title":"Session not found"}]}`), + }, nil +} + +func (*fakeTUIAPI) GetIncidentByID(context.Context, string) (*api.Incident, error) { + return nil, errors.New("incident not found") +} + +func (*fakeTUIAPI) GetAlertByID(context.Context, string) (*api.Alert, error) { + return nil, errors.New("alert not found") +} + +func (*fakeTUIAPI) APIHost() string { return "api.example.test" } + +func testSessionPage(sessions []rootly.AiChatSession, hasMore bool, before string) *rootly.ListAiChatSessionsResponse { + data := &rootly.AiChatSessionList{Sessions: sessions} + data.Meta.Pagination.HasMore = hasMore + if before == "" { + data.Meta.Pagination.NextBefore = nullable.NewNullNullable[string]() + } else { + data.Meta.Pagination.NextBefore = nullable.NewNullableWithValue(before) + } + return &rootly.ListAiChatSessionsResponse{ + HTTPResponse: &http.Response{StatusCode: http.StatusOK}, + ApplicationVndAPIJSON200: data, + } +} + +func testSkill(slug, name, description string, enabled bool, invocableBy ...string) rootly.AiSkill { + invocable := make([]rootly.AiSkillInvocableBy, 0, len(invocableBy)) + for _, value := range invocableBy { + invocable = append(invocable, rootly.AiSkillInvocableBy(value)) + } + return rootly.AiSkill{ + Slug: &slug, + Name: &name, + Description: nullable.NewNullableWithValue(description), + Enabled: &enabled, + InvocableBy: invocable, + CurrentVersion: nullable.NewNullableWithValue(3), + } +} + +func testSkillPage(skills []rootly.AiSkill, nextPage any) *rootly.ListAISkillsResponse { + data := &rootly.AiSkillList{Meta: map[string]interface{}{"next_page": nextPage}} + for _, skill := range skills { + data.Data = append(data.Data, struct { + Attributes rootly.AiSkill `json:"attributes"` + ID string `json:"id"` + Type rootly.AiSkillListDataType `json:"type"` + }{Attributes: skill, ID: skillString(skill.Slug), Type: "ai_skill"}) + } + return &rootly.ListAISkillsResponse{ + HTTPResponse: &http.Response{StatusCode: http.StatusOK}, + ApplicationVndAPIJSON200: data, + } +} + +func testHistoryPage(messages []rootly.AiChatSessionMessage, pages int) *rootly.ListAiChatSessionMessagesResponse { + return &rootly.ListAiChatSessionMessagesResponse{ + HTTPResponse: &http.Response{StatusCode: http.StatusOK}, + ApplicationVndAPIJSON200: &rootly.AiChatSessionMessageList{ + Messages: messages, + Meta: map[string]interface{}{"pagination": map[string]interface{}{"pages": pages}}, + }, + } +} + +func testSession(id, title string) rootly.AiChatSession { + return rootly.AiChatSession{ + ID: uuid.MustParse(id), Title: title, LastActiveAt: time.Now(), + } +} + +func (s *blockingStreamer) StreamAiChatEvents(ctx context.Context, _ *rootly.StreamAiChatParams) (iter.Seq2[rootly.AiChatStreamEvent, error], error) { + s.started <- struct{}{} + <-ctx.Done() + return nil, ctx.Err() +} + +func TestTUIModelAccumulatesTextChunks(t *testing.T) { + model := NewTUIModel(context.Background(), &fakeStreamer{}, TUIOptions{}) + model.state = tuiStreaming + model.stream = make(chan streamMessage, 1) + + for _, content := range []string{"first ", "second"} { + event := Event{Type: rootly.AiChatStreamEventText, Text: &rootly.AiChatStreamText{Content: content}} + _, _ = model.Update(streamMessage{event: event}) + } + if got := model.answer.String(); got != "first second" { + t.Errorf("answer = %q, want accumulated chunks", got) + } +} + +func TestTUIModelHandlesPastedText(t *testing.T) { + t.Run("idle", func(t *testing.T) { + model := NewTUIModel(context.Background(), &fakeStreamer{}, TUIOptions{}) + model.idleInterrupt = true + model.status = "press Ctrl+C again to exit" + + _, _ = model.Update(tea.PasteMsg{Content: "hello\nworld"}) + + if got := model.composer.Value(); got != "hello\nworld" { + t.Errorf("composer = %q, want pasted multiline text", got) + } + if model.idleInterrupt || model.status != "" { + t.Errorf("paste left idle interrupt state: idleInterrupt=%t status=%q", model.idleInterrupt, model.status) + } + }) + + t.Run("streaming", func(t *testing.T) { + model := NewTUIModel(context.Background(), &fakeStreamer{}, TUIOptions{}) + model.state = tuiStreaming + + _, _ = model.Update(tea.PasteMsg{Content: "hello\nworld"}) + + if got := model.composer.Value(); got != "hello\nworld" { + t.Errorf("composer = %q while streaming, want pasted multiline text", got) + } + }) + + t.Run("confirmation", func(t *testing.T) { + model := NewTUIModel(context.Background(), &fakeStreamer{}, TUIOptions{}) + model.state = tuiConfirming + model.composer.SetValue("existing draft") + + _, _ = model.Update(tea.PasteMsg{Content: "yes"}) + + if got := model.composer.Value(); got != "existing draft" { + t.Errorf("composer = %q, want existing draft unchanged", got) + } + if model.state != tuiConfirming { + t.Errorf("state = %v, want tuiConfirming", model.state) + } + }) + + t.Run("delete confirmation", func(t *testing.T) { + model := NewTUIModel(context.Background(), &fakeStreamer{}, TUIOptions{}) + model.deleteConfirming = true + model.composer.SetValue("existing draft") + + _, _ = model.Update(tea.PasteMsg{Content: "yes"}) + + if got := model.composer.Value(); got != "existing draft" { + t.Errorf("composer = %q, want existing draft unchanged", got) + } + if !model.deleteConfirming { + t.Error("paste changed delete confirmation state") + } + }) + + t.Run("session picker", func(t *testing.T) { + model := NewTUIModel(context.Background(), &fakeStreamer{}, TUIOptions{}) + model.pickerOpen = true + model.sessionPicker = list.New(nil, list.NewDefaultDelegate(), 80, 10) + model.composer.SetValue("existing draft") + + _, _ = model.Update(tea.PasteMsg{Content: "pasted text"}) + + if got := model.composer.Value(); got != "existing draft" { + t.Errorf("composer = %q, want existing draft unchanged", got) + } + }) +} + +func TestTUIModelCommitsStreamedBlocksIntoTranscript(t *testing.T) { + model := newSizedTUITestModel(80, 24) + model.state = tuiStreaming + model.turn = 1 + model.startNewTurn() + + _, cmd := model.Update(streamMessage{turn: 1, event: Event{ + Type: rootly.AiChatStreamEventText, + Text: &rootly.AiChatStreamText{Content: "First block.\n\nSecond"}, + }}) + _ = cmd + if got := ansi.Strip(model.Transcript()); strings.Count(got, "First block.") != 1 || strings.Count(got, "Second") != 1 { + t.Errorf("transcript after first chunk = %q, want committed and live text once", got) + } + + _, _ = model.Update(streamMessage{turn: 1, event: Event{ + Type: rootly.AiChatStreamEventText, + Text: &rootly.AiChatStreamText{Content: " block.\n\nTail"}, + }}) + if got := ansi.Strip(model.Transcript()); strings.Count(got, "First block.") != 1 || strings.Count(got, "Second block.") != 1 || strings.Count(got, "Tail") != 1 { + t.Errorf("transcript after second chunk = %q, want each block once", got) + } + + model.handleEvent(activityDoneEvent("complete")) + _, _ = model.Update(streamMessage{turn: 1, done: true}) + transcript := ansi.Strip(model.Transcript()) + for _, block := range []string{"First block.", "Second block.", "Tail"} { + if got := strings.Count(transcript, block); got != 1 { + t.Errorf("transcript contains %q %d times, want once: %q", block, got, transcript) + } + } + if got := model.answer.String(); got != "First block.\n\nSecond block.\n\nTail" { + t.Errorf("full answer = %q, want complete untruncated Markdown", got) + } +} + +func TestTUIStreamedMarkdownKeepsBlankLinesAcrossBlocks(t *testing.T) { + for _, test := range []struct { + name string + done bool + }{ + {name: "completed"}, + {name: "live"}, + } { + t.Run(test.name, func(t *testing.T) { + model := NewTUIModel(context.Background(), &fakeStreamer{}, TUIOptions{}) + model.Update(tea.WindowSizeMsg{Width: 80, Height: 24}) + model.state = tuiStreaming + model.turn = 1 + model.startNewTurn() + + for _, chunk := range []string{ + "Para one.\n\n", + "Para two.\n\n", + "```sh\necho hi\n\necho bye\n```\n\nAfter.", + } { + _, _ = model.Update(streamMessage{turn: 1, event: Event{ + Type: rootly.AiChatStreamEventText, + Text: &rootly.AiChatStreamText{Content: chunk}, + }}) + } + + if test.done { + model.handleEvent(activityDoneEvent("complete")) + _, _ = model.Update(streamMessage{turn: 1, done: true}) + } else if got := model.uncommittedAnswer(); got != "After." { + t.Fatalf("live remainder = %q, want After.", got) + } + + transcript := ansi.Strip(model.Transcript()) + assertStreamedMarkdownBlockSpacing(t, transcript) + if !test.done { + assertStreamedMarkdownBlockSpacing(t, ansi.Strip(model.View().Content)) + } + }) + } +} + +func assertStreamedMarkdownBlockSpacing(t *testing.T, output string) { + t.Helper() + lines := strings.Split(output, "\n") + findLine := func(fragment string, start int) int { + for index := start; index < len(lines); index++ { + if strings.Contains(strings.TrimSpace(lines[index]), fragment) { + return index + } + } + return -1 + } + + paragraphOne := findLine("● Para one.", 0) + paragraphTwo := findLine("Para two.", paragraphOne+1) + echoHi := findLine("echo hi", paragraphTwo+1) + echoBye := findLine("echo bye", echoHi+1) + after := findLine("After.", echoBye+1) + if paragraphOne < 0 || paragraphTwo < 0 || echoHi < 0 || echoBye < 0 || after < 0 { + t.Fatalf("streamed answer content missing from output: %q", output) + } + for index, want := range map[int]string{ + paragraphOne: "● Para one.", + paragraphTwo: " Para two.", + after: " After.", + } { + if got := strings.TrimRight(lines[index], " "); got != want { + t.Errorf("answer line = %q, want %q", got, want) + } + } + + for _, pair := range [][2]int{{paragraphOne, paragraphTwo}, {echoHi, echoBye}} { + if pair[1]-pair[0] != 2 || strings.TrimSpace(lines[pair[0]+1]) != "" { + t.Errorf("lines between %q and %q = %q, want exactly one blank row", lines[pair[0]], lines[pair[1]], lines[pair[0]+1:pair[1]]) + } + } + if strings.TrimSpace(lines[paragraphTwo+1]) != "" { + t.Errorf("paragraph/code separator = %q, want one blank row", lines[paragraphTwo+1]) + } + if after-echoBye < 2 || strings.TrimSpace(lines[after-1]) != "" { + t.Errorf("code/paragraph separator = %q, want one blank row before After.", lines[echoBye+1:after]) + } + + blankRun := 0 + for _, line := range lines[paragraphOne : after+1] { + if strings.TrimSpace(line) == "" { + blankRun++ + if blankRun > 1 { + t.Errorf("answer contains multiple adjacent blank rows: %q", lines[paragraphOne:after+1]) + return + } + } else { + blankRun = 0 + } + } +} + +func TestTUIStreamedFortyFiveStepsPreserveFenceAndBlankLine(t *testing.T) { + model := newSizedTUITestModel(125, 34) + model.state = tuiStreaming + model.turn = 1 + model.startNewTurn() + var answer strings.Builder + for step := 1; step <= 45; step++ { + fmt.Fprintf(&answer, "Step %d: perform the action.\n\n", step) + } + answer.WriteString("```sh\nprintf 'first command'\n\necho second command\n```\n\n") + _, _ = model.Update(streamMessage{turn: 1, event: Event{ + Type: rootly.AiChatStreamEventText, + Text: &rootly.AiChatStreamText{Content: answer.String()}, + }}) + model.handleEvent(activityDoneEvent("complete")) + _, _ = model.Update(streamMessage{turn: 1, done: true}) + transcript := ansi.Strip(model.Transcript()) + for step := 1; step <= 45; step++ { + marker := fmt.Sprintf("Step %d:", step) + if got := strings.Count(transcript, marker); got != 1 { + t.Errorf("%q appears %d times in transcript, want once", marker, got) + } + } + for _, fragment := range []string{"```sh", "printf 'first command'", "echo second command"} { + if !strings.Contains(transcript, fragment) { + t.Errorf("transcript lost fenced content %q: %q", fragment, transcript) + } + } + if !strings.Contains(transcript, "printf 'first command'\n \n echo") { + t.Errorf("transcript lost the blank line inside the fenced block: %q", transcript) + } +} + +func TestTUIFixedFooterAtCommonTerminalSizes(t *testing.T) { + for _, size := range []struct{ width, height int }{{80, 24}, {125, 34}} { + for _, blocks := range []int{0, 60} { + model := newSizedTUITestModel(size.width, size.height) + for index := 0; index < blocks; index++ { + model.println(fmt.Sprintf("transcript block %02d", index+1)) + } + lines := strings.Split(model.View().Content, "\n") + if len(lines) != size.height { + t.Fatalf("%dx%d with %d blocks rendered %d rows", size.width, size.height, blocks, len(lines)) + } + if got := ansi.Strip(lines[size.height-1]); !strings.Contains(got, "/ for commands · exit to quit") { + t.Errorf("%dx%d hint row = %q", size.width, size.height, got) + } + if got := ansi.Strip(lines[size.height-2]); !strings.Contains(got, "─") { + t.Errorf("%dx%d bottom rule = %q", size.width, size.height, got) + } + if got := ansi.Strip(lines[size.height-3]); !strings.Contains(got, "Ask Rootly") { + t.Errorf("%dx%d composer row = %q", size.width, size.height, got) + } + } + } +} + +func setTUIStreamingTestClock(model *TUIModel) { + now := time.Date(2026, time.January, 1, 12, 0, 0, 0, time.UTC) + model.now = func() time.Time { return now } + model.startedAt = now.Add(-3 * time.Second) + model.turnStartedAt = model.startedAt + model.turnVerbIndex = 0 +} + +func TestTUIStreamingStatusIsPinnedAndRemovedAfterDone(t *testing.T) { + model := newSizedTUITestModel(80, 24) + setTUIStreamingTestClock(model) + model.state = tuiStreaming + + footer := strings.Split(model.footerContent(), "\n") + if len(footer) != 6 { + t.Fatalf("streaming footer has %d rows, want 6: %q", len(footer), footer) + } + if got := ansi.Strip(footer[0]); !strings.Contains(got, "Triaging… · 3s") { + t.Errorf("streaming footer status = %q", got) + } + if footer[1] != "" { + t.Errorf("row between status and rule = %q, want blank", footer[1]) + } + if got := ansi.Strip(footer[2]); !strings.Contains(got, "─") { + t.Errorf("streaming footer top rule = %q", got) + } + if got := ansi.Strip(model.liveSection()); strings.Contains(got, "Triaging") { + t.Errorf("live section contains streaming status: %q", got) + } + if got := ansi.Strip(model.transcriptContent()); strings.Contains(got, "Triaging") { + t.Errorf("transcript content contains streaming status: %q", got) + } + if got := ansi.Strip(model.Transcript()); strings.Contains(got, "Triaging") { + t.Errorf("exit transcript contains streaming status: %q", got) + } + + model.currentPrompt = "Question" + model.answer.WriteString("Answer") + model.doneReceived = true + model.doneStatus = "complete" + _, _ = model.finishTurn(nil) + + footer = strings.Split(model.footerContent(), "\n") + if len(footer) != 4 { + t.Fatalf("completed footer has %d rows, want 4: %q", len(footer), footer) + } + if got := ansi.Strip(footer[0]); !strings.Contains(got, "─") || strings.Contains(got, "Triaging") { + t.Errorf("completed footer starts with %q, want the rule and no status", got) + } + transcript := ansi.Strip(model.Transcript()) + if !strings.Contains(transcript, "✻ Triaged for 3s") { + t.Errorf("completed transcript lost turn footer: %q", transcript) + } + if strings.Contains(transcript, "Triaging") { + t.Errorf("completed transcript contains streaming status: %q", transcript) + } +} + +func TestTUIStreamingFooterHeightPreservesViewportPosition(t *testing.T) { + bottomModel := newSizedTUITestModel(80, 24) + for index := 0; index < 60; index++ { + bottomModel.println(fmt.Sprintf("history block %02d", index+1)) + } + bottomModel.View() + if !bottomModel.viewport.AtBottom() { + t.Fatal("viewport did not start at the bottom") + } + setTUIStreamingTestClock(bottomModel) + bottomModel.state = tuiStreaming + bottomModel.View() + if !bottomModel.viewport.AtBottom() { + t.Fatal("viewport left the bottom when the streaming footer grew") + } + + scrolledModel := newSizedTUITestModel(80, 24) + for index := 0; index < 60; index++ { + scrolledModel.println(fmt.Sprintf("history block %02d", index+1)) + } + scrolledModel.View() + scrolledModel.Update(tea.KeyPressMsg{Code: tea.KeyPgUp}) + offset := scrolledModel.viewport.YOffset() + if offset == 0 { + t.Fatal("PgUp did not scroll the viewport") + } + setTUIStreamingTestClock(scrolledModel) + scrolledModel.state = tuiStreaming + scrolledModel.View() + if got := scrolledModel.viewport.YOffset(); got != offset { + t.Errorf("footer growth moved scrolled offset from %d to %d", offset, got) + } +} + +func TestTUIStreamingFooterStaysVisibleWhileScrolled(t *testing.T) { + const height = 24 + model := newSizedTUITestModel(80, height) + setTUIStreamingTestClock(model) + model.state = tuiStreaming + model.answer.WriteString(strings.Repeat("streamed answer paragraph with enough text to wrap\n\n", 40)) + model.View() + model.Update(tea.KeyPressMsg{Code: tea.KeyPgUp}) + offset := model.viewport.YOffset() + if offset == 0 { + t.Fatal("PgUp did not scroll the streamed answer") + } + + model.answer.WriteString("additional streamed answer") + view := model.View().Content + if got := model.viewport.YOffset(); got != offset { + t.Errorf("new streamed content moved scrolled offset from %d to %d", offset, got) + } + lines := strings.Split(view, "\n") + if len(lines) != height { + t.Fatalf("view rendered %d rows, want %d", len(lines), height) + } + statusRow := height - 6 + if got := ansi.Strip(lines[statusRow]); !strings.Contains(got, "Triaging… · 3s") { + t.Errorf("row %d = %q, want streaming status", statusRow, got) + } + if got := ansi.Strip(lines[height-3]); !strings.Contains(got, "Ask Rootly") { + t.Errorf("composer row = %q", got) + } + if got := ansi.Strip(model.transcriptContent()); strings.Contains(got, "Triaging") { + t.Errorf("streaming transcript contains status line: %q", got) + } +} + +func TestTUIStreamingStatusTruncatesAtNarrowWidth(t *testing.T) { + model := newSizedTUITestModel(20, 12) + setTUIStreamingTestClock(model) + model.state = tuiStreaming + model.status = strings.Repeat("streaming status ", 5) + + if got := lipgloss.Width(model.streamingStatusLine()); got != model.width-1 { + t.Errorf("streaming status width = %d, want %d", got, model.width-1) + } +} + +func TestTUIViewportFollowControls(t *testing.T) { + model := newSizedTUITestModel(80, 24) + for index := 0; index < 60; index++ { + model.println(fmt.Sprintf("history block %02d", index+1)) + } + model.View() + model.println("bottom-followed") + model.View() + if !strings.Contains(model.viewport.GetContent(), "bottom-followed") || !model.viewport.AtBottom() { + t.Fatal("viewport did not follow appended content while at bottom") + } + + model.Update(tea.KeyPressMsg{Code: tea.KeyPgUp}) + offset := model.viewport.YOffset() + if offset == 0 { + t.Fatal("PgUp did not scroll the viewport") + } + model.println("new output while scrolled") + view := ansi.Strip(model.View().Content) + if model.viewport.YOffset() != offset { + t.Errorf("append moved scrolled offset from %d to %d", offset, model.viewport.YOffset()) + } + if !strings.Contains(view, "↓ new output · ctrl+end to jump") { + t.Errorf("scrolled hint = %q, want new-output indicator", lastLine(view)) + } + + model.Update(tea.KeyPressMsg{Code: tea.KeyEnd}) + view = ansi.Strip(model.View().Content) + if !model.viewport.AtBottom() || strings.Contains(view, "↓ new output") { + t.Errorf("End did not return to bottom and clear the indicator") + } + + model.Update(tea.KeyPressMsg{Code: tea.KeyPgUp}) + if model.viewport.AtBottom() { + t.Fatal("PgUp did not scroll away from the bottom") + } + offset = model.viewport.YOffset() + model.Update(tea.KeyPressMsg{Code: tea.KeyPgDown}) + if !model.viewport.AtBottom() || model.viewport.YOffset() <= offset { + t.Fatal("PgDn did not scroll back toward the bottom") + } + offset = model.viewport.YOffset() + for _, text := range []string{"j", "k", " "} { + model.Update(tea.KeyPressMsg{Code: rune(text[0]), Text: text}) + } + if model.viewport.YOffset() != offset { + t.Errorf("typed j/k/space moved viewport from %d to %d", offset, model.viewport.YOffset()) + } + if got := model.composer.Value(); got != "jk " { + t.Errorf("composer value = %q, want typed j/k/space", got) + } + + model.composer.SetValue("end") + model.composer.CursorStart() + offset = model.viewport.YOffset() + _, _ = model.Update(tea.KeyPressMsg{Code: tea.KeyEnd}) + if model.viewport.YOffset() != offset { + t.Errorf("plain End moved viewport from %d to %d while composer had text", offset, model.viewport.YOffset()) + } + _, _ = model.Update(tea.KeyPressMsg{Code: tea.KeyBackspace}) + if got := model.composer.Value(); got != "en" { + t.Errorf("Backspace after End left composer value %q, want cursor at line end", got) + } + + _, _ = model.Update(tea.KeyPressMsg{Code: tea.KeyHome, Mod: tea.ModCtrl}) + if model.viewport.YOffset() != 0 { + t.Error("Ctrl+Home did not scroll the viewport to the top with composer text") + } + _, _ = model.Update(tea.KeyPressMsg{Code: tea.KeyEnd, Mod: tea.ModCtrl}) + if !model.viewport.AtBottom() { + t.Error("Ctrl+End did not scroll the viewport to the bottom with composer text") + } +} + +func TestTUISessionPickerConsumesPageDownWithoutScrollingViewport(t *testing.T) { + model := newSizedTUITestModel(80, 24) + for index := 0; index < 60; index++ { + model.println(fmt.Sprintf("history block %02d", index+1)) + } + model.View() + model.Update(tea.KeyPressMsg{Code: tea.KeyPgUp}) + + sessions := make([]rootly.AiChatSession, 48) + for index := range sessions { + sessions[index] = testSession(uuid.NewString(), fmt.Sprintf("Session %02d", index+1)) + } + apiClient := &fakeTUIAPI{sessionPages: []*rootly.ListAiChatSessionsResponse{ + testSessionPage(sessions, false, ""), + }} + model.api = apiClient + pageCommand := model.openSessionPicker() + result := pageCommand() + pageMessage, ok := result.(sessionPageMsg) + if !ok { + t.Fatalf("picker command result = %T, want sessionPageMsg", result) + } + _, _ = model.Update(pageMessage) + model.View() + + if model.sessionPicker.Paginator.TotalPages < 2 { + t.Fatalf("picker total pages = %d, want multiple pages", model.sessionPicker.Paginator.TotalPages) + } + page := model.sessionPicker.Paginator.Page + offset := model.viewport.YOffset() + _, _ = model.Update(tea.KeyPressMsg{Code: tea.KeyPgDown}) + if model.sessionPicker.Paginator.Page <= page { + t.Errorf("picker page = %d after PgDn, want a page after %d", model.sessionPicker.Paginator.Page, page) + } + if model.viewport.YOffset() != offset { + t.Errorf("PgDn moved viewport from %d to %d while picker was open", offset, model.viewport.YOffset()) + } +} + +func newSizedTUITestModel(width, height int) *TUIModel { + model := NewTUIModel(context.Background(), &fakeStreamer{}, TUIOptions{NoColor: true}) + model.Update(tea.WindowSizeMsg{Width: width, Height: height}) + return model +} + +func TestLastTopLevelBlockBoundaryIgnoresFencedBlankLines(t *testing.T) { + for _, marker := range []string{"```", "~~~"} { + answer := "Intro.\n\n" + marker + "go\nfirst\n\nsecond\n" + marker + "\n\nTail" + boundary := lastTopLevelBlockBoundary(answer, 0) + if got := answer[boundary:]; got != "Tail" { + t.Errorf("tail after %s fence = %q, want Tail", marker, got) + } + if block := answer[:boundary]; !strings.Contains(block, "first\n\nsecond") { + t.Errorf("committed block split %s fence contents: %q", marker, block) + } + } +} + +func TestStreamedOrderedListKeepsNumberAcrossCommittedBlocks(t *testing.T) { + model := NewTUIModel(context.Background(), &fakeStreamer{}, TUIOptions{}) + model.state = tuiStreaming + model.turn = 1 + + _, _ = model.Update(streamMessage{turn: 1, event: Event{ + Type: rootly.AiChatStreamEventText, + Text: &rootly.AiChatStreamText{Content: "1. First item\n\n"}, + }}) + _, _ = model.Update(streamMessage{turn: 1, event: Event{ + Type: rootly.AiChatStreamEventText, + Text: &rootly.AiChatStreamText{Content: "2. Second item\n\n3. Third item"}, + }}) + got := ansi.Strip(model.Transcript()) + if !strings.Contains(got, "1. First item") || !strings.Contains(got, "2. Second item") { + t.Errorf("printed list blocks lost their explicit numbering: %q", got) + } +} + +func TestCtrlCPrintsOnlyUncommittedStreamRemainder(t *testing.T) { + model := NewTUIModel(context.Background(), &fakeStreamer{}, TUIOptions{NoColor: true}) + model.state = tuiStreaming + model.answer.WriteString("Committed block\n\nUncommitted tail") + model.committedOffset = len("Committed block\n\n") + + _, cmd := model.Update(tea.KeyPressMsg{Code: 'c', Mod: tea.ModCtrl}) + _ = cmd + output := ansi.Strip(model.Transcript()) + if strings.Contains(output, "Committed block") || strings.Count(output, "Uncommitted tail") != 1 { + t.Errorf("Ctrl+C output = %q, want only the uncommitted remainder", output) + } + if strings.Count(output, "Stopped locally.") != 1 { + t.Errorf("Ctrl+C notice count in %q, want one", output) + } +} + +func TestConfirmationRequiredPrintsOnlyUncommittedStreamRemainder(t *testing.T) { + model := NewTUIModel(context.Background(), &fakeStreamer{}, TUIOptions{NoColor: true}) + model.state = tuiStreaming + model.answer.WriteString("Committed block\n\nUncommitted tail") + model.committedOffset = len("Committed block\n\n") + model.pending = &rootly.AiChatStreamConfirmationRequired{Question: "Approve?"} + model.handleEvent(activityDoneEvent("confirmation_required")) + + _, cmd := model.Update(streamMessage{turn: model.turn, done: true}) + _ = cmd + output := ansi.Strip(model.Transcript()) + if strings.Contains(output, "Committed block") || strings.Count(output, "Uncommitted tail") != 1 { + t.Errorf("confirmation output = %q, want only the uncommitted remainder", output) + } + if model.state != tuiConfirming { + t.Errorf("state = %v, want confirming", model.state) + } +} + +func TestResumedHistoryPrintsFullAssistantMessageRegardlessOfCommitOffset(t *testing.T) { + model := NewTUIModel(context.Background(), &fakeStreamer{}, TUIOptions{NoColor: true}) + model.committedOffset = len("old answer") + + model.printHistory([]rootly.AiChatSessionMessage{{ + Role: rootly.AiChatSessionMessageRole("assistant"), Content: "historical answer", + }}) + output := ansi.Strip(model.Transcript()) + if got := strings.Count(output, "historical answer"); got != 1 { + t.Errorf("history output = %q, want the complete assistant message once", output) + } +} + +func TestTUIViewClipsTallContentAboveComposer(t *testing.T) { + model := newSizedTUITestModel(40, 10) + model.status = strings.Repeat("long status content that wraps around the terminal width\n", 30) + + view := model.View().Content + if got := lipgloss.Height(view); got != 10 { + t.Fatalf("full-screen view height = %d, want 10", got) + } + lines := strings.Split(view, "\n") + if got := lines[len(lines)-1]; got != "/ for commands · exit to quit" { + t.Errorf("last row = %q, want composer hint", got) + } + if !strings.Contains(lines[len(lines)-3], "❯ ") { + t.Errorf("composer is not visible above the bottom rule: %q", view) + } +} + +func TestTUIViewBoundsLiveAnswerAndAvoidsTerminalWidthLines(t *testing.T) { + model := newSizedTUITestModel(40, 24) + model.state = tuiStreaming + model.startedAt = time.Now() + for step := 1; step <= 45; step++ { + _, _ = fmt.Fprintf(&model.answer, "%d. Step %d details %s\n\n", step, step, strings.Repeat("x", 52)) + } + _, _ = fmt.Fprintf(&model.answer, "```sh\n%s\n\n%s\n```\n45. Final marker\n", strings.Repeat("a", 80), strings.Repeat("b", 80)) + + view := model.View().Content + rendererRows := strings.Count(view, "\n") + 1 + if rendererRows != model.height { + t.Fatalf("renderer view rows = %d, want terminal height %d", rendererRows, model.height) + } + if !strings.Contains(view, "45. Final marker") { + t.Errorf("view tail = %q, want the final answer marker", view) + } + if strings.Contains(view, "1. Step 1 details") { + t.Errorf("view retained the start of the long live answer: %q", view) + } + for _, line := range strings.Split(view, "\n") { + if got := lipgloss.Width(line); got > model.width { + t.Errorf("view line width = %d, want at most terminal width %d: %q", got, model.width, line) + } + } +} + +func TestTUIRenderedAnswerLinesStayInsideTerminalWidth(t *testing.T) { + model := NewTUIModel(context.Background(), &fakeStreamer{}, TUIOptions{NoColor: true}) + model.width = 40 + + rendered := model.renderAssistantAnswerPart(strings.Repeat("x", model.width-2), true) + for _, line := range strings.Split(rendered, "\n") { + if got := lipgloss.Width(line); got >= model.width { + t.Errorf("printed answer line width = %d, want less than terminal width %d: %q", got, model.width, line) + } + } +} + +func TestTUIAnswerBulletStartsOnFirstRenderedListLine(t *testing.T) { + model := NewTUIModel(context.Background(), &fakeStreamer{}, TUIOptions{}) + model.width = 80 + + rendered := model.renderAssistantAnswerPart("1. a\n2. b", true) + lines := strings.Split(ansi.Strip(rendered), "\n") + if len(lines) == 0 || strings.TrimSpace(lines[0]) != "● 1. a" { + t.Errorf("rendered answer = %q, want first line to begin with bullet and first list item", rendered) + } +} + +func TestTUIBindingHeaderIsPrintedOnceAfterWelcomeSettles(t *testing.T) { + model := NewTUIModel(context.Background(), &fakeStreamer{}, TUIOptions{ + Incident: &api.Incident{SequentialID: "INC-8", Title: "Database latency"}, + }) + model.Update(tea.WindowSizeMsg{Width: 80, Height: 24}) + header := "INC-8 · Database latency" + if got := strings.Count(ansi.Strip(model.welcomeBanner()), header); got != 1 { + t.Errorf("header context occurs %d times, want once: %q", got, model.welcomeBanner()) + } + status := model.statusLines() + if strings.Contains(status, "Chatting about") || !strings.Contains(status, "binding INC-8: Database latency") { + t.Errorf("status = %q, want binding as a field without repeating the header", status) + } +} + +func TestTUIModelDonePrintsAndReturnsIdle(t *testing.T) { + model := NewTUIModel(context.Background(), &fakeStreamer{}, TUIOptions{}) + model.state = tuiStreaming + model.answer.WriteString("assistant reply") + model.doneStatus = "complete" + model.doneReceived = true + + updated, cmd := model.Update(streamMessage{done: true}) + if got := updated.(*TUIModel).state; got != tuiIdle { + t.Errorf("state = %v, want idle", got) + } + if cmd != nil { + t.Fatalf("done returned %T, want transcript-only update", cmd) + } + if got := ansi.Strip(model.Transcript()); !strings.Contains(got, "assistant reply") { + t.Errorf("transcript = %q, want completed answer", got) + } +} + +func TestTUIModelCtrlCCancelsStreaming(t *testing.T) { + model := NewTUIModel(context.Background(), &fakeStreamer{}, TUIOptions{}) + model.state = tuiStreaming + model.answer.WriteString("partial answer") + called := false + model.cancel = func() { called = true } + + _, cmd := model.Update(tea.KeyPressMsg{Code: 'c', Mod: tea.ModCtrl}) + if !called { + t.Error("Ctrl+C did not call the stream cancel function") + } + if model.state != tuiIdle || model.status != "" { + t.Errorf("state/status = %v/%q, want idle with no status", model.state, model.status) + } + _ = cmd + output := strings.Join(strings.Fields(ansi.Strip(model.Transcript())), " ") + if !strings.Contains(output, "partial") || !strings.Contains(output, "answer") || !strings.Contains(output, "Stopped locally. The server may still finish this turn and save it to the session.") { + t.Errorf("scrollback output = %q, want partial answer and cancellation line", output) + } +} + +func TestTUITranscriptUsesClaudeCodeStyle(t *testing.T) { + plain := renderPromptEcho("hello\nworld", true) + if plain != "❯ hello\n world" { + t.Errorf("plain prompt echo = %q, want prompt with indented continuation", plain) + } + colored := renderPromptEcho("hello\nworld", false) + if got := ansi.Strip(colored); got != "❯ hello\n world" { + t.Errorf("colored prompt echo text = %q, want prompt with indented continuation", got) + } + if strings.Contains(colored, "\x1b[48") { + t.Errorf("colored prompt echo has a background color: %q", colored) + } + model := NewTUIModel(context.Background(), &fakeStreamer{}, TUIOptions{NoColor: true}) + answer := model.renderAssistantAnswer("Hello\nworld") + if answer != "● Hello\n world" || strings.Contains(answer, "Assistant:") { + t.Errorf("assistant transcript = %q, want bullet and indented continuation without label", answer) + } +} + +func TestTUITranscriptIncludesHeaderAndConversation(t *testing.T) { + model := NewTUIModel(context.Background(), &fakeStreamer{}, TUIOptions{NoColor: true}) + model.println("❯ question") + model.println("completed answer") + model.state = tuiStreaming + model.answer.WriteString("live answer") + + transcript := model.Transcript() + for _, text := range []string{"New session", "❯ question", "completed answer", "live answer"} { + if !strings.Contains(transcript, text) { + t.Errorf("transcript = %q, missing %q", transcript, text) + } + } +} + +func TestTUITranscriptSpacingAcrossTurns(t *testing.T) { + model := newSizedTUITestModel(80, 24) + model.println("❯ first question") + model.println("● first answer") + model.printlnWithGap("✻ Triaged for 2s · done", true) + model.println("❯ second question") + model.View() + + for name, output := range map[string]string{ + "viewport": model.viewport.GetContent(), + "transcript": model.Transcript(), + } { + lines := strings.Split(strings.TrimRight(ansi.Strip(output), " \n"), "\n") + headerIndex := -1 + for index, line := range lines { + if strings.Contains(line, "New session") { + headerIndex = index + break + } + } + if headerIndex < 0 { + t.Fatalf("%s output has no session header: %q", name, output) + } + want := []string{ + "New session", + "", + "❯ first question", + "", + "● first answer", + "", + "✻ Triaged for 2s · done", + "", + "❯ second question", + } + for offset, expected := range want { + lineIndex := headerIndex + offset + if lineIndex >= len(lines) { + t.Fatalf("%s output ended before line %q: %q", name, expected, output) + } + got := strings.TrimSpace(lines[lineIndex]) + if expected == "New session" { + if !strings.HasSuffix(got, expected) { + t.Errorf("%s line %d = %q, want header ending in %q", name, lineIndex+1, got, expected) + } + } else if got != expected { + t.Errorf("%s line %d = %q, want %q", name, lineIndex+1, got, expected) + } + } + for index := 0; index+1 < len(lines); index++ { + if strings.TrimSpace(lines[index]) == "" && strings.TrimSpace(lines[index+1]) == "" { + t.Errorf("%s output has adjacent blank lines at %d: %q", name, index+1, output) + break + } + } + } +} + +func TestTUIHelpMentionsViewportNavigation(t *testing.T) { + helpText := NewTUIModel(context.Background(), &fakeStreamer{}, TUIOptions{}).helpView() + for _, text := range []string{ + "/resume, /sessions: choose a session", + "PgUp/PgDn scroll · Ctrl+Home/Ctrl+End scroll · Home/End scroll when composer is empty", + } { + if !strings.Contains(helpText, text) { + t.Errorf("help is missing %q", text) + } + } +} + +func TestTUIAnswerAndHistorySanitizeTerminalControls(t *testing.T) { + const ( + attack = "\x1b]52;c;ZXZpbA==\x07\x1b[2J\x1b]8;;https://evil.example\x07" + docsURL = "https://rootly.com/docs" + ) + model := NewTUIModel(context.Background(), &fakeStreamer{}, TUIOptions{}) + model.handleEvent(Event{Type: rootly.AiChatStreamEventText, Text: &rootly.AiChatStreamText{ + Content: "safe " + attack + " [Open docs](" + docsURL + ")", + }}) + answer := model.renderAssistantAnswer(model.answer.String()) + for _, unwanted := range []string{"]52;", "[2J", "\x1b]8;;https://evil.example"} { + if strings.Contains(answer, unwanted) { + t.Errorf("rendered answer = %q, want no %q", answer, unwanted) + } + } + if got := strings.Count(answer, ansi.SetHyperlink(docsURL)); got != 1 { + t.Errorf("answer link count = %d, want one CLI-generated OSC-8 link: %q", got, answer) + } + + historyModel := NewTUIModel(context.Background(), &fakeStreamer{}, TUIOptions{NoColor: true}) + historyModel.printHistory([]rootly.AiChatSessionMessage{ + {Role: rootly.AiChatSessionMessageRole("user"), Content: "Prompt " + attack}, + {Role: rootly.AiChatSessionMessageRole("assistant"), Content: "Answer " + attack}, + }) + history := historyModel.Transcript() + for _, unwanted := range []string{"]52;", "[2J", "\x1b"} { + if strings.Contains(history, unwanted) { + t.Errorf("printed history = %q, want no %q", history, unwanted) + } + } + historyModel.sessionTitle = "Session " + attack + status := historyModel.statusLines() + for _, unwanted := range []string{"]52;", "[2J", "\x1b"} { + if strings.Contains(status, unwanted) { + t.Errorf("status output = %q, want no %q", status, unwanted) + } + } +} + +func TestSessionPickerSanitizesTitle(t *testing.T) { + item := sessionPickerItem{session: rootly.AiChatSession{ + Title: "\x1b]52;c;ZXZpbA==\x07", + }} + if got := item.Title(); got != "Untitled" { + t.Errorf("session picker title = %q, want Untitled after sanitization", got) + } +} + +func TestTUIComposerStartsAtOneRowAndGrows(t *testing.T) { + model := NewTUIModel(context.Background(), &fakeStreamer{}, TUIOptions{}) + model.composer.SetWidth(40) + emptyRows := len(strings.Split(strings.TrimRight(model.composer.View(), "\n"), "\n")) + if emptyRows != 1 { + t.Errorf("empty composer rows = %d, want 1", emptyRows) + } + + model.composer.SetValue("one\ntwo\nthree\nfour") + filledRows := len(strings.Split(strings.TrimRight(model.composer.View(), "\n"), "\n")) + if filledRows < 3 || filledRows > model.composer.MaxHeight { + t.Errorf("multiline composer rows = %d, want between 3 and %d", filledRows, model.composer.MaxHeight) + } + _, _ = model.Update(tea.WindowSizeMsg{Width: 60, Height: 18}) + if model.composer.MaxHeight != 6 { + t.Errorf("composer max height after resize = %d, want 6", model.composer.MaxHeight) + } +} + +func TestTUIExitAliasesAndNormalExitPhrase(t *testing.T) { + for _, prompt := range []string{"exit", " Exit ", "quit"} { + model := NewTUIModel(context.Background(), &fakeStreamer{}, TUIOptions{}) + model.composer.SetValue(prompt) + cmd := model.submit() + if model.state != tuiIdle || cmd == nil { + t.Errorf("submit(%q) = state %v, command %v; want quit without streaming", prompt, model.state, cmd) + } + } + streamer := &fakeStreamer{} + model := NewTUIModel(context.Background(), streamer, TUIOptions{}) + model.composer.SetValue("exit now") + _ = model.submit() + if model.state != tuiStreaming || model.currentParams == nil || model.currentParams.Message == nil || *model.currentParams.Message != "exit now" { + t.Errorf("exit now was not submitted as a normal prompt: state=%v params=%#v", model.state, model.currentParams) + } + if model.cancel != nil { + model.cancel() + } +} + +func TestTUICommandAutocompleteTabCompletes(t *testing.T) { + model := NewTUIModel(context.Background(), &fakeStreamer{}, TUIOptions{}) + model.composer.SetValue("/res") + model.refreshCommandMenu() + if !model.commandMenuOpen { + t.Fatal("command menu did not open for a matching prefix") + } + _, _ = model.Update(tea.KeyPressMsg{Code: tea.KeyTab}) + if got := model.composer.Value(); got != "/resume " { + t.Errorf("Tab value = %q, want /resume with a trailing space", got) + } + if model.commandMenuOpen { + t.Error("command menu remained open after Tab completion") + } +} + +func TestTUICommandMenuEnterDoesNotSubmitDuringActiveTurn(t *testing.T) { + for _, test := range []struct { + name string + draft string + item slashCommandItem + }{ + {name: "new session command", draft: "/new", item: slashCommandItem{command: "/new"}}, + {name: "mention palette", draft: "@ops:resource", item: slashCommandItem{command: "@ops:resource", mention: true}}, + } { + t.Run(test.name, func(t *testing.T) { + streamer := &fakeStreamer{} + model := NewTUIModel(context.Background(), streamer, TUIOptions{}) + sessionID := uuid.MustParse(testSessionID) + model.sessionID = &sessionID + model.currentPrompt = "in-flight prompt" + model.promptHistory = []string{"in-flight prompt"} + model.currentParams = &rootly.StreamAiChatParams{} + model.turn = 7 + model.state = tuiStreaming + model.stream = make(chan streamMessage) + model.answer.WriteString("streaming answer") + model.println("persisted transcript") + model.composer.SetValue(test.draft) + model.showCommandMenu([]list.Item{test.item}) + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + model.cancel = cancel + + sessionIDBefore := model.sessionID + paramsBefore := model.currentParams + streamBefore := model.stream + transcriptBefore := model.Transcript() + _, cmd := model.Update(tea.KeyPressMsg{Code: tea.KeyEnter}) + + if cmd != nil { + t.Errorf("Enter command = %v, want nil while streaming", cmd) + } + if model.sessionID != sessionIDBefore { + t.Errorf("session id changed from %v to %v", sessionIDBefore, model.sessionID) + } + if got := model.Transcript(); got != transcriptBefore { + t.Errorf("transcript changed during active turn:\n got %q\nwant %q", got, transcriptBefore) + } + if model.stream != streamBefore || model.cancel == nil || ctx.Err() != nil { + t.Error("active stream was changed or canceled") + } + if model.currentParams != paramsBefore || model.state != tuiStreaming || model.turn != 7 { + t.Errorf("active turn changed: params=%p state=%v turn=%d", model.currentParams, model.state, model.turn) + } + if len(streamer.calls) != 0 || len(model.promptHistory) != 1 || model.currentPrompt != "in-flight prompt" { + t.Errorf("prompt was submitted: calls=%d history=%v current=%q", len(streamer.calls), model.promptHistory, model.currentPrompt) + } + if !model.commandMenuOpen || model.composer.Value() != test.draft { + t.Errorf("menu/draft changed: open=%t composer=%q", model.commandMenuOpen, model.composer.Value()) + } + }) + } +} + +func TestTUISessionsAliasOpensPicker(t *testing.T) { + for _, command := range []string{"/resume", "/sessions"} { + t.Run(command, func(t *testing.T) { + model := NewTUIModel(context.Background(), &fakeStreamer{}, TUIOptions{API: &fakeTUIAPI{}}) + cmd := model.runSlashCommand(command) + if !model.pickerOpen || !model.pickerLoading || cmd == nil { + t.Fatalf("%s: pickerOpen=%t pickerLoading=%t cmd=%v; want an opening picker and load command", command, model.pickerOpen, model.pickerLoading, cmd) + } + }) + } +} + +func TestTUICommandAutocompleteSuggestsSessionsAlias(t *testing.T) { + model := NewTUIModel(context.Background(), &fakeStreamer{}, TUIOptions{}) + model.composer.SetValue("/ses") + model.refreshCommandMenu() + if !model.commandMenuOpen { + t.Fatal("slash menu did not open for /ses") + } + for _, item := range model.commandMenu.Items() { + if command, ok := item.(slashCommandItem); ok && command.command == "/sessions" { + return + } + } + t.Errorf("slash menu items = %v, want /sessions suggestion", model.commandMenu.Items()) +} + +func TestTUISkillsPickerLoadsPagesAndHidesBuiltinSlugs(t *testing.T) { + apiClient := &fakeTUIAPI{skillPages: []*rootly.ListAISkillsResponse{ + testSkillPage([]rootly.AiSkill{ + testSkill("status", "Status skill", "reserved", true, "user"), + testSkill("incident-retro", "Incident retro", "Summarize incidents", true, "user", "model"), + }, float64(2)), + testSkillPage([]rootly.AiSkill{ + testSkill("incident-followup", "Follow-up", "Next steps", false, "model"), + }, nil), + }} + model := NewTUIModel(context.Background(), &fakeStreamer{}, TUIOptions{API: apiClient}) + cmd := model.runSlashCommand("/skills incident") + if !model.skillPickerOpen || !model.skillPickerLoading || cmd == nil { + t.Fatalf("/skills did not open a loading picker: open=%t loading=%t cmd=%v", model.skillPickerOpen, model.skillPickerLoading, cmd) + } + firstPage := cmd().(skillPickerPageMsg) + _, _ = model.Update(firstPage) + if len(apiClient.skillParams) != 1 { + t.Fatalf("skill list calls = %d, want one", len(apiClient.skillParams)) + } + if got := apiClient.skillQueries[0].Get("page[number]"); got != "1" { + t.Errorf("first page number = %q, want 1", got) + } + if got := apiClient.skillQueries[0].Get("page[size]"); got != "25" { + t.Errorf("page size = %q, want 25", got) + } + if got := apiClient.skillQueries[0].Get("sort"); got != "name" { + t.Errorf("sort = %q, want name", got) + } + if got := apiClient.skillQueries[0].Get("filter[search]"); got != "incident" { + t.Errorf("search filter = %q, want incident", got) + } + if apiClient.skillParams[0].FilterEnabled != nil { + t.Error("picker request filtered out disabled skills") + } + if len(model.skillPicker.Items()) != 2 { + t.Fatalf("picker items = %d, want skill and Load more", len(model.skillPicker.Items())) + } + first, ok := model.skillPicker.Items()[0].(skillPickerItem) + if !ok || skillString(first.skill.Slug) != "incident-retro" { + t.Fatalf("first picker item = %#v, want incident-retro (status is reserved)", model.skillPicker.Items()[0]) + } + if got := first.Title(); got != "✔ /incident-retro" { + t.Errorf("skill title = %q, want enabled marker and slug", got) + } + if got := first.Description(); got != "Incident retro · v3 · you + model · Summarize incidents" { + t.Errorf("skill description = %q", got) + } + loadMore, ok := model.skillPicker.Items()[1].(skillPickerItem) + if !ok || !loadMore.loadMore { + t.Fatalf("last picker item = %#v, want Load more", model.skillPicker.Items()[1]) + } + + _, _ = model.Update(tea.KeyPressMsg{Code: tea.KeyDown}) + _, nextPageCmd := model.Update(tea.KeyPressMsg{Code: tea.KeyEnter}) + if nextPageCmd == nil { + t.Fatal("Load more selection returned no request") + } + _, _ = model.Update(nextPageCmd()) + if len(apiClient.skillQueries) != 2 || apiClient.skillQueries[1].Get("page[number]") != "2" { + t.Fatalf("next page query = %v, want page[number]=2", apiClient.skillQueries) + } + if len(model.skillPicker.Items()) != 2 { + t.Fatalf("picker items after page 2 = %d, want two skills", len(model.skillPicker.Items())) + } + second, ok := model.skillPicker.Items()[1].(skillPickerItem) + if !ok || skillString(second.skill.Slug) != "incident-followup" { + t.Fatalf("second picker item = %#v, want appended incident-followup", model.skillPicker.Items()[1]) + } + if model.skillPickerNextPage != 0 { + t.Errorf("next page = %d, want 0 after last page", model.skillPickerNextPage) + } +} + +func TestTUISkillPickerIgnoresStaleAndCanceledPages(t *testing.T) { + apiClient := &fakeTUIAPI{skillPages: []*rootly.ListAISkillsResponse{ + testSkillPage([]rootly.AiSkill{testSkill("first", "First", "", true, "user")}, nil), + testSkillPage([]rootly.AiSkill{testSkill("second", "Second", "", true, "user")}, nil), + }} + model := NewTUIModel(context.Background(), &fakeStreamer{}, TUIOptions{API: apiClient}) + firstRequest := model.openSkillPicker("first") + stalePage := firstRequest() + _ = model.openSkillPicker("second") + _, _ = model.Update(stalePage) + if len(model.skillPicker.Items()) != 0 || !model.skillPickerLoading { + t.Fatalf("stale page changed picker: items=%v loading=%t", model.skillPicker.Items(), model.skillPickerLoading) + } + _, _ = model.Update(tea.KeyPressMsg{Code: tea.KeyEsc}) + if model.skillPickerOpen { + t.Fatal("Escape did not close the skills picker") + } +} + +func TestTUISkillPickerSelectionInsertsInvocation(t *testing.T) { + for _, test := range []struct { + name string + draft string + want string + }{ + {name: "empty draft", want: "/incident-retro "}, + {name: "existing draft", draft: " summarize INC-8", want: "/incident-retro summarize INC-8"}, + } { + t.Run(test.name, func(t *testing.T) { + apiClient := &fakeTUIAPI{skillPages: []*rootly.ListAISkillsResponse{ + testSkillPage([]rootly.AiSkill{testSkill("incident-retro", "Retro", "Summary", true, "user")}, nil), + }} + model := NewTUIModel(context.Background(), &fakeStreamer{}, TUIOptions{API: apiClient}) + pageCmd := model.openSkillPicker("")() + _, _ = model.Update(pageCmd) + model.composer.SetValue(test.draft) + _, cmd := model.Update(tea.KeyPressMsg{Code: tea.KeyEnter}) + if model.skillPickerOpen || model.composer.Value() != test.want || !model.composer.Focused() { + t.Errorf("selection state: open=%t composer=%q focused=%t, want %q and focused", model.skillPickerOpen, model.composer.Value(), model.composer.Focused(), test.want) + } + if cmd == nil { + t.Error("selection returned no composer focus command") + } + if len(apiClient.updateBodies) != 0 { + t.Errorf("skill selection made %d update calls, want none", len(apiClient.updateBodies)) + } + }) + } +} + +func TestTUISkillPickerRejectsDisabledAndModelOnlySkills(t *testing.T) { + for _, test := range []struct { + name string + skill rootly.AiSkill + message string + }{ + {name: "disabled", skill: testSkill("disabled-skill", "Disabled", "", false, "user"), message: "/disabled-skill is disabled."}, + {name: "model only", skill: testSkill("model-skill", "Model", "", true, "model"), message: "/model-skill can only be loaded by the model, not by name."}, + } { + t.Run(test.name, func(t *testing.T) { + model := NewTUIModel(context.Background(), &fakeStreamer{}, TUIOptions{NoColor: true}) + model.skillPickerOpen = true + model.skillPicker = list.New([]list.Item{skillPickerItem{skill: test.skill}}, list.NewDefaultDelegate(), 80, 5) + model.composer.SetValue("draft must remain") + _, _ = model.Update(tea.KeyPressMsg{Code: tea.KeyEnter}) + if model.skillPickerOpen || model.composer.Value() != "draft must remain" { + t.Errorf("selection changed picker/composer unexpectedly: open=%t composer=%q", model.skillPickerOpen, model.composer.Value()) + } + if !strings.Contains(model.Transcript(), test.message) { + t.Errorf("transcript = %q, want %q", model.Transcript(), test.message) + } + }) + } +} + +func TestTUISkillsPicker403And404ShowFriendlyMessage(t *testing.T) { + for _, status := range []int{http.StatusForbidden, http.StatusNotFound} { + t.Run(fmt.Sprint(status), func(t *testing.T) { + apiClient := &fakeTUIAPI{skillPages: []*rootly.ListAISkillsResponse{{ + HTTPResponse: &http.Response{StatusCode: status}, + Body: []byte(`{"errors":[{"title":"AI Skills are not enabled for this team"}]}`), + }}} + model := NewTUIModel(context.Background(), &fakeStreamer{}, TUIOptions{API: apiClient, NoColor: true}) + _, _ = model.Update(model.runSlashCommand("/skills")()) + if model.skillPickerOpen || !model.skillsUnavailable || !strings.Contains(model.Transcript(), aiSkillsDisabledMessage) { + t.Errorf("status %d state: picker=%t unavailable=%t transcript=%q", status, model.skillPickerOpen, model.skillsUnavailable, model.Transcript()) + } + model.composer.SetValue("/inc") + if cmd := model.refreshCommandMenu(); cmd != nil { + t.Errorf("status %d still scheduled skill suggestions", status) + } + }) + } +} + +func TestTUITypedSkillInvocationUsesPromptPath(t *testing.T) { + prompt := "/incident-retro summarize INC-8" + apiClient := &fakeTUIAPI{skillPages: []*rootly.ListAISkillsResponse{ + testSkillPage([]rootly.AiSkill{testSkill("incident-retro", "Retro", "", true, "user")}, nil), + }} + model := NewTUIModel(context.Background(), &fakeStreamer{}, TUIOptions{API: apiClient}) + model.composer.SetValue(prompt) + lookupCmd := model.submit() + if lookupCmd == nil { + t.Fatal("skill submission returned no lookup command") + } + _, _ = model.Update(lookupCmd()) + if model.currentParams == nil || model.currentParams.Message == nil || *model.currentParams.Message != prompt { + t.Fatalf("stream message = %#v, want the full slash prompt %q", model.currentParams, prompt) + } + if len(model.promptHistory) != 1 || model.promptHistory[0] != prompt || model.currentPrompt != prompt || model.state != tuiStreaming { + t.Errorf("prompt path state: history=%v current=%q state=%v", model.promptHistory, model.currentPrompt, model.state) + } + if len(apiClient.skillParams) != 1 || apiClient.skillParams[0].FilterSlug == nil || *apiClient.skillParams[0].FilterSlug != "incident-retro" { + t.Errorf("skill lookup params = %#v, want one lookup for incident-retro", apiClient.skillParams) + } + if model.cancel != nil { + model.cancel() + } +} + +func TestTUITypedSkillInvocationErrorsAndBuiltinPrecedence(t *testing.T) { + for _, test := range []struct { + name string + response *rootly.ListAISkillsResponse + apiError error + want string + }{ + {name: "unknown", response: testSkillPage(nil, nil), want: "Unknown command /nope. Type /help."}, + {name: "disabled", response: testSkillPage([]rootly.AiSkill{testSkill("nope", "Nope", "", false, "user")}, nil), want: "/nope is disabled."}, + {name: "model only", response: testSkillPage([]rootly.AiSkill{testSkill("nope", "Nope", "", true, "model")}, nil), want: "/nope can only be loaded by the model, not by name."}, + {name: "skills disabled", response: &rootly.ListAISkillsResponse{HTTPResponse: &http.Response{StatusCode: http.StatusForbidden}}, want: "Unknown command /nope. Type /help."}, + {name: "team key missing endpoint", response: &rootly.ListAISkillsResponse{HTTPResponse: &http.Response{StatusCode: http.StatusNotFound}}, want: "Unknown command /nope. Type /help."}, + {name: "network error", apiError: errors.New("offline\x1b[31m"), want: "Couldn't check skill /nope: offline"}, + } { + t.Run(test.name, func(t *testing.T) { + apiClient := &fakeTUIAPI{skillPages: []*rootly.ListAISkillsResponse{test.response}} + if test.apiError != nil { + apiClient.skillErrors = []error{test.apiError} + } + model := NewTUIModel(context.Background(), &fakeStreamer{}, TUIOptions{API: apiClient, NoColor: true}) + model.composer.SetValue("/nope") + cmd := model.submit() + if cmd != nil { + _, _ = model.Update(cmd()) + } + if !strings.Contains(model.Transcript(), test.want) { + t.Errorf("transcript = %q, want %q", model.Transcript(), test.want) + } + if got := model.composer.Value(); got != "/nope" { + t.Errorf("composer = %q, want the original prompt /nope", got) + } + }) + } + + apiClient := &fakeTUIAPI{} + model := NewTUIModel(context.Background(), &fakeStreamer{}, TUIOptions{API: apiClient}) + _ = model.runSlashCommand("/status") + if len(apiClient.skillParams) != 0 { + t.Errorf("/status made %d skill API calls, want none", len(apiClient.skillParams)) + } +} + +func TestTUITypedSkillInvocationFailurePreservesNewDraft(t *testing.T) { + apiClient := &fakeTUIAPI{skillErrors: []error{errors.New("offline")}} + model := NewTUIModel(context.Background(), &fakeStreamer{}, TUIOptions{API: apiClient, NoColor: true}) + model.composer.SetValue("/nope") + lookupCmd := model.submit() + model.composer.SetValue("a newer draft") + _, _ = model.Update(lookupCmd()) + if got := model.composer.Value(); got != "a newer draft" { + t.Errorf("composer = %q, want newer draft preserved", got) + } +} + +func TestTUICommandMenuSuggestsSkillsAfterBuiltins(t *testing.T) { + apiClient := &fakeTUIAPI{skillPages: []*rootly.ListAISkillsResponse{ + testSkillPage([]rootly.AiSkill{ + testSkill("incident-retro", "Incident retro", "", true, "user"), + testSkill("status", "Status skill", "", true, "user"), + }, nil), + }} + model := NewTUIModel(context.Background(), &fakeStreamer{}, TUIOptions{API: apiClient}) + model.composer.SetValue("/inc") + tickCmd := model.refreshCommandMenu() + if !model.commandMenuOpen || len(model.commandMenu.Items()) != 1 { + t.Fatalf("initial menu = %v, want built-in /incident while fetching", model.commandMenu.Items()) + } + if item := model.commandMenu.Items()[0].(slashCommandItem); item.command != "/incident" || item.skill { + t.Errorf("initial suggestion = %#v, want built-in /incident", item) + } + if tickCmd == nil { + t.Fatal("skill suggestion debounce returned no tick command") + } + _, fetchCmd := model.Update(skillSuggestionTickMsg{request: model.skillSuggestionReq, prefix: "inc"}) + if fetchCmd == nil { + t.Fatal("debounced suggestion tick returned no API request") + } + _, _ = model.Update(fetchCmd()) + if len(apiClient.skillParams) != 1 { + t.Fatalf("suggestion API calls = %d, want one", len(apiClient.skillParams)) + } + params := apiClient.skillParams[0] + if params.FilterSearch == nil || *params.FilterSearch != "inc" || + params.FilterEnabled == nil || !*params.FilterEnabled || + params.FilterInvocableBy == nil || *params.FilterInvocableBy != "user" || + apiClient.skillQueries[0].Get("page[size]") != "6" { + t.Errorf("suggestion params = %#v, query = %v", params, apiClient.skillQueries[0]) + } + var suggestions []slashCommandItem + for _, raw := range model.commandMenu.Items() { + if item, ok := raw.(slashCommandItem); ok { + suggestions = append(suggestions, item) + } + } + if len(suggestions) != 2 || suggestions[0].command != "/incident" || suggestions[0].skill || + suggestions[1].command != "/incident-retro" || !suggestions[1].skill { + t.Errorf("suggestions = %#v, want built-in /incident followed by skill /incident-retro", suggestions) + } + + model.composer.SetValue("/sk") + _ = model.refreshCommandMenu() + foundSkills := false + for _, raw := range model.commandMenu.Items() { + if item, ok := raw.(slashCommandItem); ok && item.command == "/skills" && !item.skill { + foundSkills = true + } + } + if !foundSkills { + t.Errorf("/sk menu = %v, want built-in /skills", model.commandMenu.Items()) + } + + model.skillSuggestionCache["st"] = []rootly.AiSkill{testSkill("status", "Status skill", "", true, "user")} + model.composer.SetValue("/st") + _ = model.refreshCommandMenu() + if len(model.commandMenu.Items()) != 1 { + t.Fatalf("/st menu = %v, want only built-in /status", model.commandMenu.Items()) + } + if item := model.commandMenu.Items()[0].(slashCommandItem); item.command != "/status" || item.skill { + t.Errorf("/st suggestion = %#v, want built-in /status, not a colliding skill", item) + } +} + +func TestTUICommandMenuCompletesSkillSuggestion(t *testing.T) { + for _, key := range []tea.KeyPressMsg{ + {Code: tea.KeyTab}, + {Code: tea.KeyEnter}, + } { + model := NewTUIModel(context.Background(), &fakeStreamer{}, TUIOptions{}) + model.skillSuggestionCache["incident-r"] = []rootly.AiSkill{ + testSkill("incident-retro", "Incident retro", "", true, "user"), + } + model.composer.SetValue("/incident-r") + _ = model.refreshCommandMenu() + _, _ = model.Update(key) + if got := model.composer.Value(); got != "/incident-retro " || model.commandMenuOpen { + t.Errorf("%s completion = %q, menu open=%t; want /incident-retro and closed menu", key.String(), got, model.commandMenuOpen) + } + } +} + +func TestTUISkillSuggestionResponsesAreStaleSafe(t *testing.T) { + apiClient := &fakeTUIAPI{skillPages: []*rootly.ListAISkillsResponse{ + testSkillPage([]rootly.AiSkill{testSkill("incident-retro", "Retro", "", true, "user")}, nil), + }} + model := NewTUIModel(context.Background(), &fakeStreamer{}, TUIOptions{API: apiClient}) + model.composer.SetValue("/inc") + _ = model.refreshCommandMenu() + oldRequest := model.skillSuggestionReq + oldFetch := model.loadSkillSuggestions("inc", oldRequest) + model.composer.SetValue("/incident-r") + _ = model.refreshCommandMenu() + stale := oldFetch() + _, _ = model.Update(stale) + if _, cached := model.skillSuggestionCache["inc"]; cached { + t.Fatal("stale skill suggestion response was cached") + } +} + +func TestTUIHelpListsSkills(t *testing.T) { + if got := NewTUIModel(context.Background(), &fakeStreamer{}, TUIOptions{}).helpView(); !strings.Contains(got, "/skills [query]: browse skills") { + t.Errorf("/help output does not list /skills: %q", got) + } +} + +func TestTUICommandAutocompleteEnterRunsSelectedPrefix(t *testing.T) { + apiClient := &fakeTUIAPI{sessionPages: []*rootly.ListAiChatSessionsResponse{testSessionPage(nil, false, "")}} + model := NewTUIModel(context.Background(), &fakeStreamer{}, TUIOptions{API: apiClient}) + model.composer.SetValue("/r") + model.refreshCommandMenu() + + _, cmd := model.Update(tea.KeyPressMsg{Code: tea.KeyEnter}) + if !model.pickerOpen || cmd == nil { + t.Fatalf("Enter on prefix did not run selected /resume command: picker=%v cmd=%v", model.pickerOpen, cmd) + } + result := cmd() + page, ok := result.(sessionPageMsg) + if !ok { + t.Fatalf("resume command result = %T, want session page", result) + } + _, _ = model.Update(page) + if len(apiClient.listParams) != 1 { + t.Errorf("session list calls = %d, want one selected /resume call", len(apiClient.listParams)) + } +} + +func TestTUISessionPickerLoadsMoreAndResumesSelection(t *testing.T) { + secondID := "db92b891-85c5-4e0c-a2f6-e8d873bfbde3" + apiClient := &fakeTUIAPI{sessionPages: []*rootly.ListAiChatSessionsResponse{ + testSessionPage([]rootly.AiChatSession{testSession(testSessionID, "First")}, true, "cursor-2"), + testSessionPage([]rootly.AiChatSession{testSession(secondID, "Second")}, false, ""), + }} + model := NewTUIModel(context.Background(), &fakeStreamer{}, TUIOptions{API: apiClient}) + + firstPage := model.openSessionPicker()().(sessionPageMsg) + _, _ = model.Update(firstPage) + if len(apiClient.listParams) != 1 || apiClient.listParams[0].Before != nil { + t.Fatalf("first list call params = %#v, want first page", apiClient.listParams) + } + if len(model.sessionPicker.Items()) != 2 { + t.Fatalf("picker items = %d, want session and Load more sentinel", len(model.sessionPicker.Items())) + } + + _, _ = model.Update(tea.KeyPressMsg{Code: tea.KeyDown}) + _, loadMoreCmd := model.Update(tea.KeyPressMsg{Code: tea.KeyEnter}) + if loadMoreCmd == nil { + t.Fatal("selecting Load more returned no fetch command") + } + secondPage := loadMoreCmd().(sessionPageMsg) + _, _ = model.Update(secondPage) + if len(apiClient.listParams) != 2 || apiClient.listParams[1].Before == nil || *apiClient.listParams[1].Before != "cursor-2" { + t.Fatalf("next page params = %#v, want Before cursor-2", apiClient.listParams) + } + if len(model.sessionPicker.Items()) != 2 { + t.Fatalf("picker items after more = %d, want two sessions", len(model.sessionPicker.Items())) + } + + _, _ = model.Update(tea.KeyPressMsg{Code: tea.KeyEnter}) + if model.sessionID == nil || model.sessionID.String() != secondID { + t.Errorf("selected session = %v, want %s", model.sessionID, secondID) + } + if model.sessionTitle != "Second" || model.pickerOpen { + t.Errorf("selection state title/open = %q/%v, want Second/closed", model.sessionTitle, model.pickerOpen) + } +} + +func TestTUISessionPickerSelectionRestartsWelcomeAnimation(t *testing.T) { + apiClient := &fakeTUIAPI{sessionPages: []*rootly.ListAiChatSessionsResponse{ + testSessionPage([]rootly.AiChatSession{testSession(testSessionID, "First")}, false, ""), + }} + model := NewTUIModel(context.Background(), &fakeStreamer{}, TUIOptions{API: apiClient}) + _, _ = model.Update(tea.ColorProfileMsg{Profile: colorprofile.TrueColor}) + model.Init() + settleWelcomeAnimationForTest(t, model) + + page := model.openSessionPicker()().(sessionPageMsg) + _, _ = model.Update(page) + _, cmd := model.Update(tea.KeyPressMsg{Code: tea.KeyEnter}) + if cmd == nil || !model.welcomeAnimating || model.welcomeFrame != 0 { + t.Fatalf("picker selection: cmd=%v animating=%t frame=%d; want restarted animation and command", cmd, model.welcomeAnimating, model.welcomeFrame) + } +} + +func TestTUISessionPickerEscapeCancels(t *testing.T) { + model := NewTUIModel(context.Background(), &fakeStreamer{}, TUIOptions{}) + model.pickerOpen = true + _, _ = model.Update(tea.KeyPressMsg{Code: tea.KeyEscape}) + if model.pickerOpen { + t.Error("Esc left the session picker open") + } +} + +func TestTUIHistoryRecallAndFooter(t *testing.T) { + model := NewTUIModel(context.Background(), &fakeStreamer{}, TUIOptions{NoColor: true}) + model.promptHistory = []string{"first", "second"} + _, _ = model.Update(tea.KeyPressMsg{Code: tea.KeyUp}) + if got := model.composer.Value(); got != "second" { + t.Errorf("first Up value = %q, want latest prompt", got) + } + _, _ = model.Update(tea.KeyPressMsg{Code: tea.KeyUp}) + if got := model.composer.Value(); got != "first" { + t.Errorf("second Up value = %q, want previous prompt", got) + } + _, _ = model.Update(tea.KeyPressMsg{Code: tea.KeyDown}) + if got := model.composer.Value(); got != "second" { + t.Errorf("Down value = %q, want newer prompt", got) + } + sessionID := uuid.MustParse(testSessionID) + model.sessionID = &sessionID + model.turnStartedAt = time.Now().Add(-1200 * time.Millisecond) + if footer := model.turnFooter(); !strings.Contains(footer, "✻ Triaged for ") || !strings.Contains(footer, " · done ") || strings.Contains(footer, "session ") { + t.Errorf("turn footer = %q, want elapsed without a session id", footer) + } + if status := model.statusLines(); !strings.Contains(status, "session "+testSessionID[:8]) { + t.Errorf("status = %q, want the session id to remain available", status) + } +} + +func TestTUIHistoryLoadedKeepsComposerAndSchedulesRedraw(t *testing.T) { + model := NewTUIModel(context.Background(), &fakeStreamer{}, TUIOptions{NoColor: true}) + model.historyRequest = 1 + model.Update(tea.WindowSizeMsg{Width: 80, Height: 24}) + messages := make([]rootly.AiChatSessionMessage, 0, 40) + for index := 0; index < 20; index++ { + messages = append(messages, + rootly.AiChatSessionMessage{Role: rootly.AiChatSessionMessageRole("user"), Content: fmt.Sprintf("older prompt %d", index)}, + rootly.AiChatSessionMessage{Role: rootly.AiChatSessionMessageRole("assistant"), Content: strings.Repeat("historical answer ", 40)}, + ) + } + + updated, _ := model.Update(historyLoadedMsg{ + request: 1, + title: "Resumed session", + messages: messages, + }) + if updated.(*TUIModel).state != tuiIdle { + t.Fatalf("state = %v, want idle after history loads", updated.(*TUIModel).state) + } + if view := ansi.Strip(model.View().Content); !strings.Contains(view, "❯ Ask Rootly...") || !strings.Contains(view, "/ for commands · exit to quit") { + t.Errorf("view after history load = %q, want composer visible", view) + } +} + +func TestTUIStreamingTrimsOnlyIncompleteColorLinks(t *testing.T) { + const partial = "see [Rotate](http://x?a_b" + colorModel := NewTUIModel(context.Background(), &fakeStreamer{}, TUIOptions{}) + colorModel.Update(tea.WindowSizeMsg{Width: 80, Height: 24}) + colorModel.state = tuiStreaming + colorModel.answer.WriteString(partial) + colorView := colorModel.View().Content + visible := ansi.Strip(colorView) + for _, unwanted := range []string{"http", "a_b"} { + if strings.Contains(visible, unwanted) { + t.Errorf("streaming view = %q, want no %q from incomplete link", visible, unwanted) + } + } + if final := ansi.Strip(colorModel.renderAssistantAnswer(partial)); !strings.Contains(final, "http://x?a_b") { + t.Errorf("final answer = %q, want unchanged incomplete-link text", final) + } + + rawModel := newSizedTUITestModel(80, 24) + rawModel.state = tuiStreaming + rawModel.answer.WriteString(partial) + if view := rawModel.View().Content; !strings.Contains(view, partial) { + t.Errorf("no-color streaming view = %q, want raw Markdown %q", view, partial) + } + + const completeURL = "http://x?a_b" + completeModel := NewTUIModel(context.Background(), &fakeStreamer{}, TUIOptions{}) + completeModel.Update(tea.WindowSizeMsg{Width: 80, Height: 24}) + completeModel.state = tuiStreaming + completeModel.answer.WriteString("see [Rotate](" + completeURL + ") and more") + completeView := completeModel.View().Content + completeVisible := ansi.Strip(completeView) + if !strings.Contains(completeVisible, "Rotate") || !strings.Contains(completeVisible, "and more") { + t.Errorf("complete-link streaming view = %q, want link label and trailing text", completeVisible) + } + if got := strings.Count(completeView, ansi.SetHyperlink(completeURL)); got != 1 { + t.Errorf("complete-link count = %d, want one hyperlink", got) + } +} + +func TestLoadSessionHistoryFetchesLastPage(t *testing.T) { + firstMessage := rootly.AiChatSessionMessage{ + ID: uuid.MustParse(testSessionID), + Role: rootly.AiChatSessionMessageRole("user"), Content: "old prompt", + CreatedAt: time.Date(2026, 1, 1, 0, 0, 0, 0, time.UTC), + } + lastMessage := rootly.AiChatSessionMessage{ + ID: uuid.MustParse("c532fe53-20b9-49ee-ab28-68e917d09e73"), + Role: rootly.AiChatSessionMessageRole("assistant"), Content: "latest answer", + CreatedAt: time.Date(2026, 1, 2, 0, 0, 0, 0, time.UTC), + } + apiClient := &fakeTUIAPI{messagePages: []*rootly.ListAiChatSessionMessagesResponse{ + testHistoryPage([]rootly.AiChatSessionMessage{firstMessage}, 2), + testHistoryPage([]rootly.AiChatSessionMessage{lastMessage}, 2), + }} + + messages, _, err := loadSessionHistory(context.Background(), apiClient, uuid.MustParse(testSessionID)) + if err != nil { + t.Fatalf("loadSessionHistory() error = %v", err) + } + if len(apiClient.messageParams) != 2 { + t.Fatalf("history calls = %d, want first and last page", len(apiClient.messageParams)) + } + if *apiClient.messageParams[0].PageNumber != 1 || *apiClient.messageParams[0].PageSize != 100 || + *apiClient.messageParams[1].PageNumber != 2 || *apiClient.messageParams[1].PageSize != 100 { + t.Errorf("history page params = %#v, %#v; want pages 1 and 2 with size 100", apiClient.messageParams[0], apiClient.messageParams[1]) + } + if len(messages) != 2 || messages[0].Content != "old prompt" || messages[1].Content != "latest answer" { + t.Errorf("history messages = %#v, want first and last page messages in order", messages) + } +} + +func TestTUIResumeByIDLoadsTitleAndPrintsIt(t *testing.T) { + apiClient := &fakeTUIAPI{messagePages: []*rootly.ListAiChatSessionMessagesResponse{ + { + ApplicationVndAPIJSON200: &rootly.AiChatSessionMessageList{Title: "Incident review\x1b[2J"}, + }, + }} + sessionID := uuid.MustParse(testSessionID) + model := NewTUIModel(context.Background(), &fakeStreamer{}, TUIOptions{ + NoColor: true, + API: apiClient, + SessionID: &sessionID, + }) + initSequence := reflect.ValueOf(model.Init()()) + if initSequence.Kind() != reflect.Slice || initSequence.Len() != 2 { + t.Fatalf("initial resume commands = %v, want composer focus and history load", initSequence) + } + + history := model.loadHistory(sessionID)().(historyLoadedMsg) + updated, _ := model.Update(history) + resumed := updated.(*TUIModel) + if resumed.sessionTitle != "Incident review" { + t.Fatalf("session title = %q, want sanitized history title", resumed.sessionTitle) + } + if output := ansi.Strip(resumed.Transcript()); !strings.Contains(output, "Incident review") { + t.Errorf("resume transcript = %q, want the loaded session title in the header", output) + } + if status := resumed.statusLines(); !strings.Contains(status, "title Incident review") { + t.Errorf("status lines = %q, want loaded session title", status) + } + if got := strings.Count(ansi.Strip(resumed.Transcript()), "Incident review"); got != 1 { + t.Errorf("resumed title appears %d times in transcript, want one header context", got) + } +} + +func TestTUIResumeByIDFallsBackToShortIDWithoutTitle(t *testing.T) { + for _, test := range []struct { + name string + title string + }{ + {name: "absent"}, + {name: "untitled", title: "Untitled"}, + } { + t.Run(test.name, func(t *testing.T) { + apiClient := &fakeTUIAPI{messagePages: []*rootly.ListAiChatSessionMessagesResponse{ + { + ApplicationVndAPIJSON200: &rootly.AiChatSessionMessageList{Title: test.title}, + }, + }} + sessionID := uuid.MustParse(testSessionID) + model := NewTUIModel(context.Background(), &fakeStreamer{}, TUIOptions{ + NoColor: true, + API: apiClient, + SessionID: &sessionID, + }) + + history := model.loadHistory(sessionID)().(historyLoadedMsg) + updated, _ := model.Update(history) + resumed := updated.(*TUIModel) + if resumed.sessionTitle != "" { + t.Fatalf("session title = %q, want empty fallback state", resumed.sessionTitle) + } + if got, want := resumed.sessionHeaderContext(), shortSessionID(testSessionID); got != want { + t.Errorf("resume header context = %q, want %q", got, want) + } + }) + } +} + +func TestTUITaskActivityIsTerminalAndDeduplicated(t *testing.T) { + model := NewTUIModel(context.Background(), &fakeStreamer{}, TUIOptions{NoColor: true}) + model.handleEvent(Event{Type: rootly.AiChatStreamEventTaskUpdate, TaskUpdate: &rootly.AiChatStreamTaskUpdate{ID: "task-1", Title: "Create incident", Status: "in_progress"}}) + model.handleEvent(Event{Type: rootly.AiChatStreamEventTaskUpdate, TaskUpdate: &rootly.AiChatStreamTaskUpdate{ID: "task-1", Title: "Create incident", Status: "complete"}}) + model.handleEvent(Event{Type: rootly.AiChatStreamEventTaskUpdate, TaskUpdate: &rootly.AiChatStreamTaskUpdate{ID: "task-1", Title: "Create incident", Status: "complete"}}) + model.handleEvent(Event{Type: rootly.AiChatStreamEventTaskUpdate, TaskUpdate: &rootly.AiChatStreamTaskUpdate{ID: "task-2", Title: "Notify owner", Status: "cancelled"}}) + model.handleEvent(Event{Type: rootly.AiChatStreamEventTaskUpdate, TaskUpdate: &rootly.AiChatStreamTaskUpdate{ID: "task-3", Title: "Searched incidents (1.2s)", Status: "complete"}}) + got := ansi.Strip(model.Transcript()) + if strings.Count(got, "✓ Create incident") != 1 || strings.Count(got, "✗ Notify owner") != 1 { + t.Errorf("task lines = %q; want one success and one failure glyph", got) + } + if !strings.Contains(got, "✓ Searched incidents (1.2s)") || strings.Contains(got, "(1.2s) (") { + t.Errorf("duration task line = %q, want title duration only", got) + } +} + +func TestTUIStatusActivityIsPersisted(t *testing.T) { + model := NewTUIModel(context.Background(), &fakeStreamer{}, TUIOptions{NoColor: true}) + model.handleEvent(activityStatusEvent("Get incidents...")) + model.handleEvent(activityStatusEvent("Composing response...")) + model.handleEvent(Event{Type: rootly.AiChatStreamEventText, Text: &rootly.AiChatStreamText{Content: "Answer"}}) + model.handleEvent(activityDoneEvent("complete")) + if got := ansi.Strip(model.Transcript()); !strings.Contains(got, "✓ Get incidents (") { + t.Errorf("activity transcript = %q, want successful Get incidents", got) + } +} + +func TestTUIFinishTurnTreatsConfirmedActivityAsSuccessful(t *testing.T) { + model := NewTUIModel(context.Background(), &fakeStreamer{}, TUIOptions{NoColor: true}) + model.activities.Handle(activityStatusEvent("Get incidents...")) + model.doneStatus = "confirmed" + model.doneReceived = true + + _, _ = model.finishTurn(nil) + if output := ansi.Strip(model.Transcript()); !strings.Contains(output, "✓ Get incidents (") { + t.Errorf("finish output = %q, want successful activity line", output) + } +} + +func TestTUIConfirmationDefersTextFinalizedActivityUntilResumedTurn(t *testing.T) { + model := NewTUIModel(context.Background(), &fakeStreamer{}, TUIOptions{NoColor: true}) + model.state = tuiStreaming + model.handleEvent(activityStatusEvent("Create action item...")) + model.handleEvent(activityStatusEvent("Composing response...")) + model.handleEvent(Event{Type: rootly.AiChatStreamEventText, Text: &rootly.AiChatStreamText{Content: "Ready."}}) + model.handleEvent(Event{ + Type: rootly.AiChatStreamEventConfirmationRequired, + ConfirmationRequired: &rootly.AiChatStreamConfirmationRequired{ + CheckpointID: testCheckpointID, + Question: "Approve?", + ExpiresAt: "2099-01-01T00:00:00Z", + }, + }) + model.handleEvent(activityDoneEvent("confirmation_required")) + + _, _ = model.finishTurn(nil) + if output := ansi.Strip(model.Transcript()); strings.Contains(output, "✓ Create action item (") { + t.Fatalf("confirmation output = %q, want no activity before approval", output) + } + if model.state != tuiConfirming { + t.Fatalf("state = %v, want confirming", model.state) + } + + model.pending = nil + model.state = tuiStreaming + model.answer.Reset() + model.doneStatus = "" + model.handleEvent(activityStatusEvent("Create action item...")) + model.handleEvent(Event{Type: rootly.AiChatStreamEventText, Text: &rootly.AiChatStreamText{Content: "Created."}}) + model.handleEvent(activityDoneEvent("complete")) + + _, _ = model.finishTurn(nil) + if output := ansi.Strip(model.Transcript()); strings.Count(output, "✓ Create action item (") != 1 { + t.Fatalf("resumed output = %q, want exactly one completed activity", output) + } +} + +func TestTUIConfirmationFlushesStatusFinalization(t *testing.T) { + model := NewTUIModel(context.Background(), &fakeStreamer{}, TUIOptions{NoColor: true}) + model.state = tuiStreaming + model.handleEvent(activityStatusEvent("A...")) + model.handleEvent(activityStatusEvent("B...")) + model.handleEvent(Event{ + Type: rootly.AiChatStreamEventConfirmationRequired, + ConfirmationRequired: &rootly.AiChatStreamConfirmationRequired{ + CheckpointID: testCheckpointID, + Question: "Approve?", + ExpiresAt: "2099-01-01T00:00:00Z", + }, + }) + model.handleEvent(activityDoneEvent("confirmation_required")) + + _, _ = model.finishTurn(nil) + output := ansi.Strip(model.Transcript()) + if strings.Count(output, "✓ A (") != 1 || strings.Contains(output, "✓ B (") { + t.Fatalf("confirmation output = %q, want retained A line and no B line", output) + } +} + +func TestTUIFinishTurnAppendsAnswerToTranscript(t *testing.T) { + model := NewTUIModel(context.Background(), &fakeStreamer{}, TUIOptions{NoColor: true}) + model.println(renderPromptEcho("Question", model.noColor)) + model.answer.WriteString("Answer") + model.currentPrompt = "Question" + model.doneStatus = "complete" + model.doneReceived = true + + _, cmd := model.finishTurn(nil) + if cmd != nil { + t.Fatalf("finishTurn() command = %T, want no inline print command", cmd) + } + if got := ansi.Strip(model.Transcript()); !strings.Contains(got, "Answer") || !strings.Contains(got, "Question") { + t.Errorf("transcript = %q, want prompt and answer", got) + } +} + +func TestTUISessionCommandsRequireSessionAndMapErrors(t *testing.T) { + apiClient := &fakeTUIAPI{} + model := NewTUIModel(context.Background(), &fakeStreamer{}, TUIOptions{API: apiClient}) + _ = model.runSlashCommand("/rename name") + if output := ansi.Strip(model.Transcript()); !strings.Contains(output, "No session yet: send a message first.") { + t.Errorf("no-session output = %q, want exact no-session message", output) + } + if len(apiClient.updateBodies) != 0 { + t.Fatalf("updates = %d, want no API call without a session", len(apiClient.updateBodies)) + } + + sessionID := uuid.MustParse(testSessionID) + model.sessionID = &sessionID + apiClient.updateStatus = http.StatusForbidden + cmd := model.runSlashCommand("/pin") + result, ok := cmd().(sessionActionMsg) + if !ok || result.err == nil || !strings.Contains(result.err.Error(), "ai.chat:write") { + t.Fatalf("pin result = %#v, want mapped API permission error", result) + } + if len(apiClient.updateBodies) != 1 { + t.Errorf("updates = %d, want one PATCH", len(apiClient.updateBodies)) + } +} + +func TestTUISessionActionsMapNotFoundForBoundSessions(t *testing.T) { + const want = "Session not found. Incident and alert chats can't be renamed, pinned, archived or deleted from the CLI yet." + apiClient := &fakeTUIAPI{updateStatus: http.StatusNotFound, deleteStatus: http.StatusNotFound} + model := NewTUIModel(context.Background(), &fakeStreamer{}, TUIOptions{API: apiClient}) + sessionID := uuid.MustParse(testSessionID) + model.sessionID = &sessionID + + for _, command := range []string{"/rename Updated title", "/pin", "/unpin", "/archive"} { + t.Run(command, func(t *testing.T) { + cmd := model.runSlashCommand(command) + if cmd == nil { + t.Fatal("session action returned no command") + } + msg := cmd() + result, ok := msg.(sessionActionMsg) + if !ok { + t.Fatalf("command result = %T, want sessionActionMsg", msg) + } + if result.err == nil { + t.Fatal("session action error = nil, want not-found error") + } + _ = model.handleSessionAction(result) + if output := strings.Join(strings.Fields(ansi.Strip(model.Transcript())), " "); !strings.Contains(output, want) { + t.Errorf("session action output = %q, want %q", output, want) + } + }) + } + + model.runSlashCommand("/delete") + _, cmd := model.Update(tea.KeyPressMsg{Code: 'y', Text: "y"}) + if cmd == nil { + t.Fatal("delete confirmation returned no API command") + } + result, ok := cmd().(sessionActionMsg) + if !ok || result.err == nil { + t.Fatalf("delete result = %#v, want not-found sessionActionMsg", result) + } + _ = model.handleSessionAction(result) + if output := strings.Join(strings.Fields(ansi.Strip(model.Transcript())), " "); !strings.Contains(output, want) { + t.Errorf("delete output = %q, want %q", output, want) + } +} + +func TestTUIBindingCommandsShowUsageWithoutID(t *testing.T) { + model := NewTUIModel(context.Background(), &fakeStreamer{}, TUIOptions{}) + for _, test := range []struct { + command string + want string + }{ + {command: "/incident", want: "Usage: /incident "}, + {command: "/alert", want: "Usage: /alert "}, + } { + t.Run(test.command, func(t *testing.T) { + _ = model.runSlashCommand(test.command) + if output := ansi.Strip(model.Transcript()); !strings.Contains(output, test.want) { + t.Errorf("output = %q, want %q", output, test.want) + } + }) + } +} + +func TestTUIExportWritesTranscriptWithoutOverwriting(t *testing.T) { + dir := t.TempDir() + path := filepath.Join(dir, "chat.md") + sessionID := uuid.MustParse(testSessionID) + model := NewTUIModel(context.Background(), &fakeStreamer{}, TUIOptions{}) + model.sessionID = &sessionID + model.conversation = []transcriptTurn{{user: "question", assistant: "answer"}} + + if cmd := model.exportTranscript(path); cmd != nil { + t.Fatalf("export command = %T, want transcript-only update", cmd) + } + got, err := os.ReadFile(path) + if err != nil { + t.Fatalf("read exported file: %v", err) + } + if !strings.Contains(string(got), "## You\n\nquestion") || !strings.Contains(string(got), "## Rootly\n\nanswer") { + t.Errorf("export = %q, want both transcript sections", got) + } + if cmd := model.exportTranscript(path); cmd != nil { + t.Fatalf("second export command = %T, want transcript-only update", cmd) + } + gotAgain, err := os.ReadFile(path) + if err != nil { + t.Fatalf("read exported file after repeated export: %v", err) + } + if !bytes.Equal(gotAgain, got) { + t.Errorf("existing export changed: before %q, after %q", got, gotAgain) + } +} + +func TestTUIModelDropsStaleStreamMessages(t *testing.T) { + streamer := &blockingStreamer{started: make(chan struct{}, 2)} + model := NewTUIModel(context.Background(), streamer, TUIOptions{}) + + model.startStream(&rootly.StreamAiChatParams{}) + firstTurn := model.turn + <-streamer.started + _, _ = model.Update(tea.KeyPressMsg{Code: 'c', Mod: tea.ModCtrl}) + + model.startStream(&rootly.StreamAiChatParams{}) + secondTurn := model.turn + t.Cleanup(func() { + if model.cancel != nil { + model.cancel() + } + }) + <-streamer.started + + _, _ = model.Update(streamMessage{turn: firstTurn, done: true}) + if model.state != tuiStreaming || model.turn != secondTurn { + t.Errorf("state/turn = %v/%d, want streaming on turn %d", model.state, model.turn, secondTurn) + } +} + +func TestTUIModelCtrlCTwiceWhileIdleQuits(t *testing.T) { + model := NewTUIModel(context.Background(), &fakeStreamer{}, TUIOptions{}) + _, firstCmd := model.Update(tea.KeyPressMsg{Code: 'c', Mod: tea.ModCtrl}) + if firstCmd != nil || model.status != "press Ctrl+C again to exit" { + t.Fatalf("first Ctrl+C = %v, %q, want status message and no quit", firstCmd, model.status) + } + _, secondCmd := model.Update(tea.KeyPressMsg{Code: 'c', Mod: tea.ModCtrl}) + if secondCmd == nil { + t.Fatal("second Ctrl+C did not return tea.Quit") + } + if _, ok := secondCmd().(tea.QuitMsg); !ok { + t.Errorf("second Ctrl+C command result = %T, want tea.QuitMsg", secondCmd()) + } + if view := model.View().Content; view != "" { + t.Errorf("quitting view = %q, want empty view", view) + } +} + +func TestTUIModelConfirmStartsResume(t *testing.T) { + model := NewTUIModel(context.Background(), &fakeStreamer{}, TUIOptions{}) + sessionID := uuid.MustParse(testSessionID) + model.sessionID = &sessionID + model.state = tuiConfirming + model.pending = &rootly.AiChatStreamConfirmationRequired{ + CheckpointID: testCheckpointID, + Question: "Approve?", + ExpiresAt: "2099-01-01T00:00:00Z", + } + + _, cmd := model.Update(tea.KeyPressMsg{Code: 'y', Text: "y"}) + if cmd == nil || model.state != tuiStreaming { + t.Fatalf("confirm key returned %v with state %v, want streaming resume", cmd, model.state) + } + if model.currentParams == nil || model.currentParams.Decision == nil || *model.currentParams.Decision != rootly.StreamAiChatParamsDecisionConfirm { + t.Errorf("resume params = %#v, want decision confirm", model.currentParams) + } + if model.currentParams.Message != nil { + t.Errorf("resume message = %v, want nil", model.currentParams.Message) + } +} + +func TestTUIModelExpiredConfirmationDoesNotResume(t *testing.T) { + model := NewTUIModel(context.Background(), &fakeStreamer{}, TUIOptions{}) + sessionID := uuid.MustParse(testSessionID) + model.sessionID = &sessionID + model.state = tuiConfirming + model.pending = &rootly.AiChatStreamConfirmationRequired{ + CheckpointID: testCheckpointID, + Question: "Approve?", + ExpiresAt: "2000-01-01T00:00:00Z", + } + + _, cmd := model.Update(tea.KeyPressMsg{Code: 'y', Text: "y"}) + if model.state != tuiIdle || model.pending != nil || model.currentParams != nil { + t.Errorf("expired confirmation state = %v, pending=%v, params=%v; want idle without resume", model.state, model.pending, model.currentParams) + } + _ = cmd + if output := ansi.Strip(model.Transcript()); !strings.Contains(output, "Confirmation expired; no action was taken.") { + t.Errorf("scrollback output = %q, want expiration message", output) + } +} + +func TestTUIModelExpiredCancelDoesNotResume(t *testing.T) { + streamer := &fakeStreamer{} + model := NewTUIModel(context.Background(), streamer, TUIOptions{}) + sessionID := uuid.MustParse(testSessionID) + model.sessionID = &sessionID + model.state = tuiConfirming + model.pending = &rootly.AiChatStreamConfirmationRequired{ + CheckpointID: testCheckpointID, + Question: "Approve?", + ExpiresAt: "2000-01-01T00:00:00Z", + } + + _, cmd := model.Update(tea.KeyPressMsg{Code: 'n', Text: "n"}) + if model.state != tuiIdle || model.pending != nil || len(streamer.calls) != 0 { + t.Errorf("expired cancel state = %v, pending=%v, calls=%d; want idle without resume", model.state, model.pending, len(streamer.calls)) + } + _ = cmd + if output := ansi.Strip(model.Transcript()); !strings.Contains(output, "Confirmation expired; no action was taken.") { + t.Errorf("scrollback output = %q, want expiration message", output) + } +} + +func TestTUIModelSeparatesResumedText(t *testing.T) { + streamer := &fakeStreamer{responses: [][]rootly.AiChatStreamEvent{{}}} + model := NewTUIModel(context.Background(), streamer, TUIOptions{}) + sessionID := uuid.MustParse(testSessionID) + model.sessionID = &sessionID + model.state = tuiConfirming + model.pending = &rootly.AiChatStreamConfirmationRequired{ + CheckpointID: testCheckpointID, + ExpiresAt: "2099-01-01T00:00:00Z", + } + model.answer.WriteString("A") + + _ = model.handleConfirmationDecision(rootly.StreamAiChatParamsDecisionConfirm) + model.handleEvent(Event{Type: rootly.AiChatStreamEventText, Text: &rootly.AiChatStreamText{Content: "B"}}) + if got := model.answer.String(); got != "A\n\nB" { + t.Errorf("answer = %q, want resumed text separated by a blank line", got) + } + if model.cancel != nil { + model.cancel() + } +} + +func TestTUITranscriptSeparatesPromptAnswerAndFooter(t *testing.T) { + now := time.Date(2026, time.January, 1, 12, 0, 0, 0, time.UTC) + model := NewTUIModel(context.Background(), &fakeStreamer{}, TUIOptions{NoColor: true}) + model.Update(tea.WindowSizeMsg{Width: 80, Height: 24}) + model.now = func() time.Time { return now } + model.currentPrompt = "hi" + model.println(renderPromptEcho("hi", true)) + model.state = tuiStreaming + model.startedAt = now + model.turn = 1 + model.turnStartedAt = now.Add(-8 * time.Second) + model.turnVerbIndex = 0 + model.answer.WriteString("hello") + + const promptAndAnswer = "❯ hi\n\n● hello" + if transcript := ansi.Strip(model.Transcript()); !strings.Contains(transcript, promptAndAnswer) { + t.Errorf("live transcript = %q, want one blank line between prompt and answer", transcript) + } + viewLines := strings.Split(ansi.Strip(model.View().Content), "\n") + for index := range viewLines { + viewLines[index] = strings.TrimRight(viewLines[index], " ") + } + if view := strings.Join(viewLines, "\n"); !strings.Contains(view, promptAndAnswer) { + t.Errorf("live view = %q, want one blank line between prompt and answer", view) + } + + model.doneReceived = true + model.doneStatus = "complete" + _, _ = model.finishTurn(nil) + const completeTurn = "❯ hi\n\n● hello\n\n✻ Triaged for 8s · done 12:00 PM" + if transcript := ansi.Strip(model.Transcript()); !strings.Contains(transcript, completeTurn) { + t.Errorf("completed transcript = %q, want one blank line before the footer", transcript) + } + + history := NewTUIModel(context.Background(), &fakeStreamer{}, TUIOptions{NoColor: true}) + history.printHistory([]rootly.AiChatSessionMessage{ + {Role: rootly.AiChatSessionMessageRole("user"), Content: "hi"}, + {Role: rootly.AiChatSessionMessageRole("assistant"), Content: "hello"}, + {Role: rootly.AiChatSessionMessageRole("user"), Content: "next"}, + }) + if transcript := ansi.Strip(history.Transcript()); !strings.Contains(transcript, "❯ hi\n\n● hello\n\n❯ next") { + t.Errorf("resumed history = %q, want one blank line between turns", transcript) + } +} + +func TestTUIHistoryMarkdownReflowsAfterResize(t *testing.T) { + const answer = "The response should explain that incident severity can move from SEV0 (critical, widespread) down to lower levels as responders contain impact, restore affected systems, and verify recovery with stakeholders.\n\nThe second paragraph adds more detail: teams should keep ownership clear, communicate updates, and document decisions made during recovery." + model := NewTUIModel(context.Background(), &fakeStreamer{}, TUIOptions{}) + model.Update(tea.WindowSizeMsg{Width: 139, Height: 32}) + model.printHistory([]rootly.AiChatSessionMessage{ + {Role: rootly.AiChatSessionMessageRole("user"), Content: "Summarize the incident."}, + {Role: rootly.AiChatSessionMessageRole("assistant"), Content: answer}, + }) + + for _, width := range []int{139, 110, 60} { + if width != 139 { + model.Update(tea.WindowSizeMsg{Width: width, Height: 32}) + } + model.View() + transcript := ansi.Strip(model.Transcript()) + assertTranscriptFitsWithoutRepeatedBlankLines(t, "transcript", transcript, width) + assertTranscriptFitsWithoutRepeatedBlankLines(t, "viewport", ansi.Strip(model.viewport.GetContent()), width) + if width == 110 { + for _, line := range strings.Split(transcript, "\n") { + if strings.Contains(line, "(critical,") && !strings.Contains(line, "widespread)") { + t.Errorf("110-column transcript breaks after (critical,: %q", line) + } + } + } + } +} + +func TestTUIStreamedMarkdownReflowsAfterResize(t *testing.T) { + const answer = "The response should explain that incident severity can move from SEV0 (critical, widespread) down to lower levels as responders contain impact, restore affected systems, and verify recovery with stakeholders.\n\nThe second paragraph adds more detail: teams should keep ownership clear, communicate updates, and document decisions made during recovery." + now := time.Date(2026, time.January, 1, 12, 0, 0, 0, time.UTC) + model := NewTUIModel(context.Background(), &fakeStreamer{}, TUIOptions{}) + model.Update(tea.WindowSizeMsg{Width: 139, Height: 32}) + model.now = func() time.Time { return now } + model.currentPrompt = "Summarize the incident." + model.println(renderPromptEcho(model.currentPrompt, model.noColor)) + model.state = tuiStreaming + model.turn = 1 + model.turnStartedAt = now.Add(-4 * time.Second) + model.startedAt = now.Add(-4 * time.Second) + + _, _ = model.Update(streamMessage{ + turn: model.turn, + event: Event{Type: rootly.AiChatStreamEventText, Text: &rootly.AiChatStreamText{Content: answer}}, + }) + if model.committedOffset == 0 { + t.Fatal("streamed answer did not commit its first paragraph") + } + model.handleEvent(activityDoneEvent("complete")) + _, _ = model.Update(streamMessage{turn: model.turn, done: true}) + + model.View() + assertTranscriptFitsWithoutRepeatedBlankLines(t, "139-column transcript", ansi.Strip(model.Transcript()), 139) + model.Update(tea.WindowSizeMsg{Width: 110, Height: 28}) + model.View() + transcript := ansi.Strip(model.Transcript()) + assertTranscriptFitsWithoutRepeatedBlankLines(t, "resized transcript", transcript, 110) + assertTranscriptFitsWithoutRepeatedBlankLines(t, "resized viewport", ansi.Strip(model.viewport.GetContent()), 110) + for _, line := range strings.Split(transcript, "\n") { + if strings.Contains(line, "(critical,") && !strings.Contains(line, "widespread)") { + t.Errorf("resized streamed transcript breaks after (critical,: %q", line) + } + } +} + +func TestTUITranscriptMarkdownCacheRendersOncePerWidth(t *testing.T) { + model := NewTUIModel(context.Background(), &fakeStreamer{}, TUIOptions{}) + model.Update(tea.WindowSizeMsg{Width: 139, Height: 24}) + model.printHistory([]rootly.AiChatSessionMessage{{ + Role: rootly.AiChatSessionMessageRole("assistant"), Content: "A Markdown answer to cache.", + }}) + model.Transcript() + initialRenderCount := model.renderCount + if initialRenderCount != 1 { + t.Fatalf("initial render count = %d, want one", initialRenderCount) + } + + model.println("A plain appended line.") + model.Transcript() + if model.renderCount != initialRenderCount { + t.Errorf("render count after plain append = %d, want %d", model.renderCount, initialRenderCount) + } + + model.Update(tea.WindowSizeMsg{Width: 110, Height: 24}) + model.Transcript() + resizedRenderCount := model.renderCount + if resizedRenderCount != initialRenderCount+1 { + t.Errorf("render count after resize = %d, want %d", resizedRenderCount, initialRenderCount+1) + } + model.Transcript() + if model.renderCount != resizedRenderCount { + t.Errorf("render count after unchanged rebuild = %d, want %d", model.renderCount, resizedRenderCount) + } +} + +func assertTranscriptFitsWithoutRepeatedBlankLines(t *testing.T, name, output string, width int) { + t.Helper() + lines := strings.Split(output, "\n") + previousLineBlank := false + for index, line := range lines { + if lineWidth := ansi.StringWidth(line); lineWidth > width { + t.Errorf("%s line %d has width %d, exceeds %d: %q", name, index+1, lineWidth, width, line) + } + lineBlank := strings.TrimSpace(line) == "" + if lineBlank && previousLineBlank { + t.Errorf("%s has adjacent blank lines at %d: %q", name, index+1, output) + break + } + previousLineBlank = lineBlank + } +} + +func TestTUIPickersRenderFullScreenAndRestoreChat(t *testing.T) { + tests := []struct { + name string + title string + open func(*TUIModel) tea.Cmd + isOpen func(*TUIModel) bool + }{ + { + name: "sessions", + title: "Sessions", + open: func(model *TUIModel) tea.Cmd { + return model.openSessionPicker() + }, + isOpen: func(model *TUIModel) bool { return model.pickerOpen }, + }, + { + name: "skills", + title: "AI Skills", + open: func(model *TUIModel) tea.Cmd { + return model.openSkillPicker("") + }, + isOpen: func(model *TUIModel) bool { return model.skillPickerOpen }, + }, + } + for _, test := range tests { + t.Run(test.name, func(t *testing.T) { + model := NewTUIModel(context.Background(), &fakeStreamer{}, TUIOptions{ + API: &fakeTUIAPI{}, NoColor: true, + }) + model.Update(tea.WindowSizeMsg{Width: 80, Height: 24}) + model.println("chat transcript marker") + for index := 0; index < 30; index++ { + model.println(fmt.Sprintf("history %d", index)) + } + model.composer.SetValue("composer draft") + model.state = tuiStreaming + model.status = "stream status marker" + model.View() + model.viewport.GotoTop() + offset := model.viewport.YOffset() + + if cmd := test.open(model); cmd == nil { + t.Fatal("opening picker returned no load command") + } + picker := model.View() + if picker.MouseMode != tea.MouseModeNone { + t.Errorf("picker mouse mode = %v, want none", picker.MouseMode) + } + pickerView := picker.Content + plainPickerView := ansi.Strip(pickerView) + if strings.Contains(plainPickerView, "chat transcript marker") || + strings.Contains(plainPickerView, "composer draft") || + strings.Contains(plainPickerView, "stream status marker") { + t.Errorf("picker view contains chat content: %q", plainPickerView) + } + if !strings.Contains(plainPickerView, test.title) { + t.Errorf("picker view = %q, want title %q", plainPickerView, test.title) + } + if lines := len(strings.Split(strings.TrimRight(pickerView, "\n"), "\n")); lines != 24 { + t.Errorf("picker view has %d rows, want 24", lines) + } + + _, _ = model.Update(tea.WindowSizeMsg{Width: 100, Height: 30}) + if lines := len(strings.Split(strings.TrimRight(model.View().Content, "\n"), "\n")); lines != 30 { + t.Errorf("resized picker view has %d rows, want 30", lines) + } + _, _ = model.Update(tea.KeyPressMsg{Code: tea.KeyEscape}) + if test.isOpen(model) { + t.Fatal("Escape did not close the picker") + } + chatView := ansi.Strip(model.View().Content) + if !strings.Contains(chatView, "chat transcript marker") || + !strings.Contains(chatView, "composer draft") || + !strings.Contains(chatView, "stream status marker") { + t.Errorf("chat view was not restored: %q", chatView) + } + if got := model.composer.Value(); got != "composer draft" { + t.Errorf("composer value = %q, want draft preserved", got) + } + if got := model.viewport.YOffset(); got != offset { + t.Errorf("viewport offset = %d, want unchanged %d", got, offset) + } + }) + } +} + +func TestTUIModelStreamErrorsPrintPartialAnswer(t *testing.T) { + tests := []struct { + name string + message streamMessage + doneStatus string + doneReceived bool + wantMessage string + }{ + { + name: "error event", + message: streamMessage{event: Event{ + Type: rootly.AiChatStreamEventError, + Error: &rootly.AiChatStreamError{Message: "event failed"}, + }}, + wantMessage: "event failed", + }, + { + name: "done error", + message: streamMessage{done: true}, + doneStatus: "error", + doneReceived: true, + wantMessage: "AI chat turn failed", + }, + { + name: "truncated stream", + message: streamMessage{done: true}, + wantMessage: "AI chat stream ended before the turn finished", + }, + { + name: "transport error", + message: streamMessage{done: true, err: errors.New("transport failed")}, + wantMessage: "transport failed", + }, + } + for _, test := range tests { + t.Run(test.name, func(t *testing.T) { + model := NewTUIModel(context.Background(), &fakeStreamer{}, TUIOptions{}) + model.state = tuiStreaming + model.doneStatus = test.doneStatus + model.doneReceived = test.doneReceived + model.answer.WriteString("partial answer") + + _, _ = model.Update(test.message) + output := ansi.Strip(model.Transcript()) + if !strings.Contains(output, "partial") || !strings.Contains(output, "answer") || !strings.Contains(output, test.wantMessage) { + t.Errorf("scrollback output = %q, want partial answer and %q", output, test.wantMessage) + } + if model.state != tuiIdle { + t.Errorf("state = %v, want idle", model.state) + } + }) + } +} + +func TestTUIModelNoColorErrorHasNoANSI(t *testing.T) { + model := NewTUIModel(context.Background(), &fakeStreamer{}, TUIOptions{NoColor: true}) + model.state = tuiStreaming + model.answer.WriteString("# partial answer") + + _, _ = model.Update(streamMessage{done: true, err: errors.New("stream failed")}) + output := model.Transcript() + if strings.Contains(output, "\x1b[") { + t.Errorf("scrollback output contains ANSI escape sequence: %q", output) + } + if !strings.Contains(output, "# partial answer") || !strings.Contains(output, "stream failed") { + t.Errorf("scrollback output = %q, want raw Markdown and error", output) + } +} + +func TestTUIModelCachesMarkdownRenderingAcrossViews(t *testing.T) { + model := NewTUIModel(context.Background(), &fakeStreamer{}, TUIOptions{}) + model.Update(tea.WindowSizeMsg{Width: 80, Height: 24}) + model.state = tuiStreaming + model.answer.WriteString("A streamed answer") + + model.View() + model.View() + if model.renderCount != 1 { + t.Errorf("render count = %d, want one render for unchanged answer", model.renderCount) + } +} + +func TestTUIModelResumeStatuses(t *testing.T) { + tests := []struct { + status string + wantMessage string + wantStatus string + }{ + {status: "cancelled", wantMessage: "Action cancelled."}, + {status: "expired", wantMessage: "Confirmation expired; no action was taken."}, + {status: "confirmed"}, + {status: "pending", wantStatus: "turn halted: pending"}, + } + for _, test := range tests { + t.Run(test.status, func(t *testing.T) { + model := NewTUIModel(context.Background(), &fakeStreamer{}, TUIOptions{}) + model.state = tuiStreaming + model.doneStatus = test.status + model.doneReceived = true + model.answer.WriteString("reply") + + _, _ = model.Update(streamMessage{done: true}) + output := ansi.Strip(model.Transcript()) + if test.wantMessage != "" && !strings.Contains(output, test.wantMessage) { + t.Errorf("scrollback output = %q, want %q", output, test.wantMessage) + } + if test.status == "confirmed" && !strings.Contains(output, "reply") { + t.Errorf("confirmed output = %q, want completed answer", output) + } + if model.status != test.wantStatus { + t.Errorf("status = %q, want %q", model.status, test.wantStatus) + } + }) + } +} + +func TestTUIModelNewCommandClearsSession(t *testing.T) { + model := NewTUIModel(context.Background(), &fakeStreamer{}, TUIOptions{}) + sessionID := uuid.MustParse(testSessionID) + model.sessionID = &sessionID + model.resumedSession = true + model.sessionTitle = "Old session" + model.println("old history") + model.conversation = []transcriptTurn{{user: "old question", assistant: "old answer"}} + model.composer.SetValue("/new") + model.showCommandMenu([]list.Item{slashCommandItem{command: "/new"}}) + + _, _ = model.Update(tea.KeyPressMsg{Code: tea.KeyEnter}) + if model.sessionID != nil { + t.Errorf("session id = %v, want nil", model.sessionID) + } + if got := ansi.Strip(model.Transcript()); strings.Contains(got, "old history") || strings.Contains(got, "Old session") || !strings.Contains(got, "New session") { + t.Errorf("new-session transcript = %q, want only a fresh header", got) + } + if len(model.conversation) != 0 { + t.Errorf("conversation after /new = %#v, want empty", model.conversation) + } +} + +func TestTUIModelModifiedEnterInsertsNewline(t *testing.T) { + for _, modifier := range []tea.KeyMod{tea.ModShift, tea.ModAlt} { + model := NewTUIModel(context.Background(), &fakeStreamer{}, TUIOptions{}) + model.composer.SetValue("question") + + _, _ = model.Update(tea.KeyPressMsg{Code: tea.KeyEnter, Mod: modifier}) + if model.state != tuiIdle { + t.Errorf("modifier %v submitted the prompt", modifier) + } + if !strings.Contains(model.composer.Value(), "\n") { + t.Errorf("modifier %v did not insert a newline: %q", modifier, model.composer.Value()) + } + } +} diff --git a/internal/agent/verbs.go b/internal/agent/verbs.go new file mode 100644 index 0000000..897c393 --- /dev/null +++ b/internal/agent/verbs.go @@ -0,0 +1,19 @@ +package agent + +var statusVerbs = []struct { + Active string + Done string +}{ + {Active: "Triaging", Done: "Triaged"}, + {Active: "Investigating", Done: "Investigated"}, + {Active: "Correlating", Done: "Correlated"}, + {Active: "Untangling", Done: "Untangled"}, + {Active: "Spelunking", Done: "Spelunked"}, + {Active: "Sleuthing", Done: "Sleuthed"}, + {Active: "Wrangling", Done: "Wrangled"}, + {Active: "Deciphering", Done: "Deciphered"}, + {Active: "Escalating thoughts", Done: "Escalated thoughts"}, + {Active: "Paging the hamsters", Done: "Paged the hamsters"}, + {Active: "Rummaging through logs", Done: "Rummaged through logs"}, + {Active: "Reticulating runbooks", Done: "Reticulated runbooks"}, +} diff --git a/internal/agent/welcome.go b/internal/agent/welcome.go new file mode 100644 index 0000000..4d7a31b --- /dev/null +++ b/internal/agent/welcome.go @@ -0,0 +1,194 @@ +package agent + +import ( + "os" + "path/filepath" + "strings" + + "charm.land/lipgloss/v2" + "github.com/charmbracelet/x/ansi" + rootly "github.com/rootlyhq/rootly-go" +) + +const ( + welcomeLogoWidth = 14 + welcomeLogoHeight = 4 +) + +var welcomeLogoLines = []string{ + " ▄█ ", + " ▄▀ █ ", + "▀▀▀▀▀ █ ▄▀▀▀▀", + " ▀█ ", +} + +func (m *TUIModel) welcomeBanner() string { + return m.welcomeBannerContent(pulseLogo(m.noColor)) + "\n" +} + +func (m *TUIModel) welcomeBannerFrame() string { + return m.welcomeBannerContent(renderPulseFrame(m.welcomeFrame)) +} + +func (m *TUIModel) welcomeBannerContent(logo string) string { + text := m.welcomeTextRows() + if m.width < 50 { + return text + } + textWidth := max(0, m.width-welcomeLogoWidth-2) + lines := strings.Split(text, "\n") + for index, line := range lines { + lines[index] = ansi.Truncate(line, textWidth, "") + } + return lipgloss.JoinHorizontal(lipgloss.Center, logo, " ", strings.Join(lines, "\n")) +} + +func (m *TUIModel) welcomeTextRows() string { + version := strings.TrimPrefix(sanitizeTerminalText(m.version), "v") + versionText := "" + if version != "" && version != "dev" { + versionText = " v" + version + } + title := sanitizeTerminalText("✦") + " " + sanitizeTerminalText("Rootly AI") + sanitizeTerminalText(versionText) + if !m.noColor { + title = lipgloss.NewStyle().Foreground(lipgloss.Color("#F5A24F")).Render(sanitizeTerminalText("✦")) + + " " + lipgloss.NewStyle().Bold(true).Render(sanitizeTerminalText("Rootly AI")) + if versionText != "" { + title += faint(false, sanitizeTerminalText(versionText)) + } + } + + userLine := "" + if m.currentUser != nil { + attributes := m.currentUser.Data.Attributes + name := currentUserFirstName(attributes) + team := currentUserTeamName(attributes) + if name == "" { + name = sanitizeTerminalText(attributes.Email) + } + if name != "" { + name = "Welcome back, " + name + } + if name != "" && team != "" { + name += " · " + team + } + userLine = name + } + + host := m.apiHost + if host == "" && m.api != nil { + host = m.api.APIHost() + } + hostLine := host + contextLine := m.sessionHeaderContext() + lines := []string{title, userLine} + if dir := displayWorkingDir(m.workingDir); dir != "" { + lines = append(lines, dir) + } + lines = append(lines, hostLine) + if capability := m.welcomeCapabilityLine(); capability != "" { + lines = append(lines, capability) + } + lines = append(lines, contextLine) + for index := 1; index < len(lines); index++ { + line := lines[index] + line = sanitizeTerminalText(line) + if !m.noColor { + line = faint(false, line) + } + lines[index] = line + } + return strings.Join(lines, "\n") +} + +func (m *TUIModel) sessionHeaderContext() string { + if m.incident != nil { + id := m.incident.SequentialID + if id == "" { + id = shortSessionID(m.incident.ID) + } + return id + " · " + sanitizeTerminalText(m.incident.Title) + } + if m.alert != nil { + id := m.alert.ShortID + if id == "" { + id = shortSessionID(m.alert.ID) + } + return id + " · " + sanitizeTerminalText(m.alert.Summary) + } + if m.resumedSession { + if title := normalizeSessionTitle(m.sessionTitle); title != "" { + return title + } + if m.sessionID != nil { + return shortSessionID(m.sessionID.String()) + } + } + return "New session" +} + +func pulseLogo(noColor bool) string { + logo := strings.Join(welcomeLogoLines, "\n") + if noColor { + return logo + } + return lipgloss.NewStyle().Foreground(lipgloss.Color("#7C5CE6")).Render(logo) +} + +const welcomeCapabilityText = "Ask me about incidents, alerts and on-call" + +func (m *TUIModel) welcomeCapabilityLine() string { + available := m.width + if m.width >= 50 { + available = m.width - welcomeLogoWidth - 2 + } + if available <= 0 || ansi.StringWidth(welcomeCapabilityText) > available { + return "" + } + return welcomeCapabilityText +} + +func displayWorkingDir(dir string) string { + if dir == "" { + return "" + } + if home, err := os.UserHomeDir(); err == nil && home != "" { + if dir == home { + return "~" + } + if rest, ok := strings.CutPrefix(dir, home+string(os.PathSeparator)); ok { + return filepath.Join("~", rest) + } + } + return dir +} + +func currentUserFirstName(attributes rootly.CurrentUser) string { + if first := strings.TrimSpace(attributes.FirstName.GetOrEmpty()); first != "" { + return sanitizeTerminalText(first) + } + if fields := strings.Fields(currentUserName(attributes)); len(fields) > 0 { + return fields[0] + } + return "" +} + +func currentUserName(attributes rootly.CurrentUser) string { + name := attributes.FullName.GetOrEmpty() + if name == "" { + name = strings.TrimSpace(attributes.FirstName.GetOrEmpty() + " " + attributes.LastName.GetOrEmpty()) + } + return sanitizeTerminalText(name) +} + +func currentUserTeamName(attributes rootly.CurrentUser) string { + fullNameWithTeam := sanitizeTerminalText(attributes.FullNameWithTeam.GetOrEmpty()) + if !strings.HasPrefix(fullNameWithTeam, "[") { + return "" + } + end := strings.Index(fullNameWithTeam, "] ") + if end <= 1 { + return "" + } + return strings.TrimSpace(fullNameWithTeam[1:end]) +} diff --git a/internal/agent/welcome_animation.go b/internal/agent/welcome_animation.go new file mode 100644 index 0000000..9bdd1e1 --- /dev/null +++ b/internal/agent/welcome_animation.go @@ -0,0 +1,118 @@ +package agent + +import ( + "math" + "strings" + "time" + + tea "charm.land/bubbletea/v2" + "charm.land/lipgloss/v2" + "github.com/charmbracelet/colorprofile" +) + +const ( + welcomeDrawSpeed = 1.2 + welcomeBannerFrameInterval = 45 * time.Millisecond + welcomeGlimmerCenter = 3 +) + +var ( + welcomeCrest = []string{"#FFFFFF", "#E0DAFB", "#C9BEF7", "#9D86F0", "#7C5CE6"} + welcomeGlimmer = []string{ + "#9D86F0", "#C9BEF7", "#E0DAFB", "#FFFFFF", "#E0DAFB", "#C9BEF7", "#9D86F0", + } + welcomeDrawFrames = int(math.Ceil((welcomeDrawMaxDistance() + float64(len(welcomeCrest))) / welcomeDrawSpeed)) + welcomeSettleFrame = welcomeDrawFrames + welcomeGlimmerFrames() + 4 +) + +type welcomeBannerTickMsg struct { + generation int +} + +func (m *TUIModel) welcomeBannerTick() tea.Cmd { + generation := m.welcomeGeneration + return tea.Tick(welcomeBannerFrameInterval, func(time.Time) tea.Msg { + return welcomeBannerTickMsg{generation: generation} + }) +} + +func (m *TUIModel) restartWelcomeAnimation() tea.Cmd { + if m.noColor || !m.colorProfileKnown || m.colorProfile < colorprofile.ANSI256 { + m.welcomeAnimating = false + return nil + } + m.welcomeGeneration++ + m.welcomeFrame = 0 + m.welcomeAnimating = true + return m.welcomeBannerTick() +} + +func welcomeDrawMaxDistance() float64 { + maxDistance := 0.0 + for row, line := range welcomeLogoLines { + for column, character := range []rune(line) { + if character == ' ' { + continue + } + if distance := welcomeDrawDistance(row, column); distance > maxDistance { + maxDistance = distance + } + } + } + return maxDistance +} + +func welcomeDrawDistance(_, column int) float64 { + return float64(column) +} + +func welcomeCellColor(frame, row, column int) string { + if frame >= welcomeSettleFrame { + return "#7C5CE6" + } + if frame > welcomeDrawFrames && frame <= welcomeDrawFrames+welcomeGlimmerFrames() { + phase := (frame - welcomeDrawFrames - 1) * 2 + index := column - phase + welcomeGlimmerCenter + if index >= 0 && index < len(welcomeGlimmer) { + return welcomeGlimmer[index] + } + return "#7C5CE6" + } + if frame > welcomeDrawFrames { + return "#7C5CE6" + } + edge := float64(frame)*welcomeDrawSpeed - welcomeDrawDistance(row, column) + if edge < 0 { + return "" + } + return welcomeCrest[min(int(math.Floor(edge)), len(welcomeCrest)-1)] +} + +func welcomeGlimmerPeriod() int { + return welcomeLogoWidth + welcomeLogoHeight + 16 +} + +func welcomeGlimmerFrames() int { + return (welcomeGlimmerPeriod() + 1) / 2 +} + +func renderPulseFrame(frame int) string { + lines := make([]string, len(welcomeLogoLines)) + for row, line := range welcomeLogoLines { + var rendered strings.Builder + for column, character := range []rune(line) { + if character == ' ' { + rendered.WriteRune(character) + continue + } + color := welcomeCellColor(frame, row, column) + if color == "" { + rendered.WriteRune(' ') + continue + } + rendered.WriteString(lipgloss.NewStyle().Foreground(lipgloss.Color(color)).Render(string(character))) + } + lines[row] = rendered.String() + } + return strings.Join(lines, "\n") +} diff --git a/internal/agent/welcome_test.go b/internal/agent/welcome_test.go new file mode 100644 index 0000000..3b15fca --- /dev/null +++ b/internal/agent/welcome_test.go @@ -0,0 +1,452 @@ +package agent + +import ( + "context" + "strings" + "testing" + "time" + + tea "charm.land/bubbletea/v2" + "charm.land/lipgloss/v2" + "github.com/charmbracelet/colorprofile" + "github.com/charmbracelet/x/ansi" + "github.com/google/uuid" + "github.com/oapi-codegen/nullable" + rootly "github.com/rootlyhq/rootly-go" + + "github.com/rootlyhq/rootly-cli/internal/api" +) + +func TestWelcomeBannerIncludesVersionUserTeamAndHost(t *testing.T) { + model := NewTUIModel(context.Background(), &fakeStreamer{}, TUIOptions{ + NoColor: true, + Version: "v1.2.3", + APIHost: "api.rootly.com", + CurrentUser: currentUserForWelcome( + "Avery Rootly", + "avery@example.com", + "[Incident Response] Avery Rootly", + ), + }) + + banner := ansi.Strip(model.welcomeBanner()) + for _, expected := range []string{ + " ▄█ ", + " ▄▀ █ ", + "▀▀▀▀▀ █ ▄▀▀▀▀", + " ▀█ ", + "Rootly AI", + "v1.2.3", + "Welcome back, Avery · Incident Response", + "api.rootly.com", + "New session", + } { + if !strings.Contains(banner, expected) { + t.Errorf("welcome banner = %q, want %q", banner, expected) + } + } + if strings.Contains(banner, "vv1.2.3") { + t.Errorf("welcome banner = %q, want one leading v", banner) + } + bannerLines := strings.Split(strings.TrimSuffix(banner, "\n"), "\n") + textRows := strings.Split(model.welcomeTextRows(), "\n") + if len(bannerLines) != max(welcomeLogoHeight, len(textRows)) { + t.Errorf("welcome banner rows = %d, want %d", len(bannerLines), max(welcomeLogoHeight, len(textRows))) + } + logoStart := -1 + for index, line := range bannerLines { + if strings.HasPrefix(line, welcomeLogoLines[0]) { + logoStart = index + break + } + } + if logoStart < 0 || logoStart+len(welcomeLogoLines) > len(bannerLines) { + t.Fatalf("welcome banner = %q, want the logo rows", banner) + } + for index, expected := range welcomeLogoLines { + if !strings.HasPrefix(bannerLines[logoStart+index], expected) { + t.Errorf("welcome banner row %d = %q, want logo row %q at its start", logoStart+index+1, bannerLines[logoStart+index], expected) + } + } + for index, line := range strings.Split(pulseLogo(true), "\n") { + if width := lipgloss.Width(line); width != welcomeLogoWidth { + t.Errorf("pulse row %d width = %d, want %d", index+1, width, welcomeLogoWidth) + } + } + model.Init() + if model.welcomeAnimating { + t.Error("no-color welcome banner should be static") + } +} + +func TestWelcomeBannerNoColorAndSanitizesUser(t *testing.T) { + model := NewTUIModel(context.Background(), &fakeStreamer{}, TUIOptions{ + NoColor: true, + Version: "dev", + APIHost: "api.rootly.com", + CurrentUser: currentUserForWelcome( + "Avery \x1b]52;c;ZXZpbA==\x07\x1b[2JRootly", + "avery@example.com", + "", + ), + }) + + textRows := model.welcomeTextRows() + if strings.Contains(textRows, "\x1b") { + t.Errorf("welcome text rows = %q, want no ANSI escapes in no-color mode", textRows) + } + + banner := model.welcomeBanner() + for _, unwanted := range []string{"\x1b", "]52;", "[2J"} { + if strings.Contains(banner, unwanted) { + t.Errorf("welcome banner = %q, want no %q", banner, unwanted) + } + } + if strings.Contains(banner, "vdev") { + t.Errorf("welcome banner = %q, want development version omitted", banner) + } + model.Init() + if strings.Contains(banner, "\x1b") || model.welcomeAnimating { + t.Errorf("no-color welcome banner = %q, animating = %t", banner, model.welcomeAnimating) + } + if !strings.Contains(banner, "Welcome back, Avery") { + t.Errorf("welcome banner = %q, want safe user name text", banner) + } +} + +func TestWelcomeTitleStyling(t *testing.T) { + model := NewTUIModel(context.Background(), &fakeStreamer{}, TUIOptions{Version: "1.2.3"}) + title := strings.Split(model.welcomeTextRows(), "\n")[0] + if got := ansi.Strip(title); got != "✦ Rootly AI v1.2.3" { + t.Fatalf("plain title = %q, want title and version", got) + } + + orange := strings.Index(title, "\x1b[38;2;245;162;79m") + bold := strings.Index(title, "\x1b[1m") + faint := strings.Index(title, "\x1b[2m") + if orange < 0 || bold < 0 { + t.Fatalf("title = %q, want orange and bold SGR styling", title) + } + if faint < 0 || faint < orange || faint < bold { + t.Errorf("title = %q, want only the version faint after the accented title", title) + } +} + +func TestWelcomeBannerAnimationSettlesInPlace(t *testing.T) { + model := NewTUIModel(context.Background(), &fakeStreamer{}, TUIOptions{Version: "1.2.3"}) + animationDuration := time.Duration(welcomeSettleFrame) * welcomeBannerFrameInterval + if animationDuration <= 0 || animationDuration > 2200*time.Millisecond { + t.Fatalf("welcome animation duration = %s, want >0 and <=2.2s", animationDuration) + } + _, _ = model.Update(tea.ColorProfileMsg{Profile: colorprofile.TrueColor}) + model.Init() + if !model.welcomeAnimating { + t.Fatal("welcome animation did not start") + } + _, _ = model.Update(tea.WindowSizeMsg{Width: 80, Height: 24}) + view := ansi.Strip(model.View().Content) + if !strings.Contains(view, "Rootly AI") || strings.HasPrefix(view, "\n") { + t.Errorf("animated view = %q, want in-place header and no leading padding", view) + } + + for frame := 0; frame < welcomeSettleFrame-1; frame++ { + _, _ = model.Update(welcomeBannerTickMsg{}) + if !model.welcomeAnimating { + t.Fatalf("animation settled early at frame %d", frame+1) + } + } + _, _ = model.Update(welcomeBannerTickMsg{}) + if model.welcomeAnimating { + t.Fatal("animation did not settle at the expected frame") + } + _, _ = model.Update(welcomeBannerTickMsg{}) + if model.welcomeAnimating { + t.Error("an extra tick restarted the animation") + } +} + +func TestWelcomeAnimationRestartsForNewSessionCommands(t *testing.T) { + for _, command := range []string{"/new", "/clear"} { + t.Run(command, func(t *testing.T) { + model := NewTUIModel(context.Background(), &fakeStreamer{}, TUIOptions{}) + _, _ = model.Update(tea.ColorProfileMsg{Profile: colorprofile.TrueColor}) + model.Init() + settleWelcomeAnimationForTest(t, model) + + cmd := model.runSlashCommand(command) + if !model.welcomeAnimating || model.welcomeFrame != 0 || cmd == nil { + t.Fatalf("after %s: animating=%t frame=%d cmd=%v; want restarted animation and tick", command, model.welcomeAnimating, model.welcomeFrame, cmd) + } + for range welcomeSettleFrame { + _, _ = model.Update(welcomeBannerTickMsg{generation: model.welcomeGeneration}) + } + if model.welcomeAnimating { + t.Errorf("animation did not settle after %d current-generation ticks", welcomeSettleFrame) + } + }) + } +} + +func TestWelcomeAnimationRestartIgnoresStaleTick(t *testing.T) { + model := NewTUIModel(context.Background(), &fakeStreamer{}, TUIOptions{}) + _, _ = model.Update(tea.ColorProfileMsg{Profile: colorprofile.TrueColor}) + model.Init() + if !model.welcomeAnimating { + t.Fatal("welcome animation did not start") + } + oldGeneration := model.welcomeGeneration + + if cmd := model.restartWelcomeAnimation(); cmd == nil { + t.Fatal("restart returned no tick command") + } + if model.welcomeGeneration == oldGeneration || model.welcomeFrame != 0 || !model.welcomeAnimating { + t.Fatalf("restart state: generation=%d frame=%d animating=%t", model.welcomeGeneration, model.welcomeFrame, model.welcomeAnimating) + } + _, _ = model.Update(welcomeBannerTickMsg{generation: oldGeneration}) + if model.welcomeFrame != 0 { + t.Errorf("stale tick advanced frame to %d, want 0", model.welcomeFrame) + } +} + +func TestWelcomeAnimationDoesNotRestartWithoutColor(t *testing.T) { + model := NewTUIModel(context.Background(), &fakeStreamer{}, TUIOptions{NoColor: true}) + model.Init() + if cmd := model.runSlashCommand("/new"); cmd != nil { + t.Errorf("no-color /new command = %v, want no tick", cmd) + } + if model.welcomeAnimating { + t.Error("no-color /new started the welcome animation") + } +} + +func settleWelcomeAnimationForTest(t *testing.T, model *TUIModel) { + t.Helper() + for frame := 0; frame < welcomeSettleFrame; frame++ { + _, _ = model.Update(welcomeBannerTickMsg{generation: model.welcomeGeneration}) + } + if model.welcomeAnimating { + t.Fatalf("welcome animation remained active after %d ticks", welcomeSettleFrame) + } +} + +func TestWelcomeGlimmerReachesLogoEdges(t *testing.T) { + topLeftColumn := -1 + for column, character := range []rune(welcomeLogoLines[0]) { + if character != ' ' { + topLeftColumn = column + break + } + } + bottomRow := len(welcomeLogoLines) - 1 + bottomRightColumn := -1 + for column, character := range []rune(welcomeLogoLines[bottomRow]) { + if character != ' ' { + bottomRightColumn = column + } + } + if topLeftColumn < 0 || bottomRightColumn < 0 { + t.Fatal("welcome logo has no edge cells") + } + + for _, cell := range []struct { + name string + row int + column int + }{ + {name: "top-left", row: 0, column: topLeftColumn}, + {name: "bottom-right", row: bottomRow, column: bottomRightColumn}, + } { + t.Run(cell.name, func(t *testing.T) { + for frame := welcomeDrawFrames + 1; frame <= welcomeDrawFrames+welcomeGlimmerFrames(); frame++ { + if color := welcomeCellColor(frame, cell.row, cell.column); color != "" && color != "#7C5CE6" { + return + } + } + t.Errorf("cell (%d,%d) never gets a glimmer color", cell.row, cell.column) + }) + } +} + +func TestWelcomePulseDrawsLeftToRight(t *testing.T) { + midDrawFrame := welcomeDrawFrames / 2 + if color := welcomeCellColor(midDrawFrame, 2, 0); color == "" { + t.Error("left edge of pulse is still blank midway through drawing") + } + if color := welcomeCellColor(midDrawFrame, 2, 13); color != "" { + t.Errorf("right edge of pulse has color %q midway through drawing, want blank", color) + } + for _, column := range []int{0, 13} { + if color := welcomeCellColor(welcomeSettleFrame, 2, column); color != "#7C5CE6" { + t.Errorf("settled cell (2,%d) color = %q, want #7C5CE6", column, color) + } + } +} + +func TestWelcomeGlimmerStartsAtLeftEdge(t *testing.T) { + for frame := welcomeDrawFrames + 1; frame <= welcomeDrawFrames+welcomeGlimmerFrames(); frame++ { + color := welcomeCellColor(frame, 2, 0) + if color == "" || color == "#7C5CE6" { + continue + } + if rightColor := welcomeCellColor(frame, 2, 13); rightColor != "#7C5CE6" { + t.Errorf("right edge glimmer color = %q at frame %d, want settled purple while beat starts left", rightColor, frame) + } + return + } + t.Fatal("no glimmer frame lights the left edge with a non-settled color") +} + +func TestEarlyPromptSettlesWelcomeBannerBeforePromptEcho(t *testing.T) { + model := NewTUIModel(context.Background(), &fakeStreamer{responses: [][]rootly.AiChatStreamEvent{{}}}, TUIOptions{}) + model.width = 100 + model.Init() + _, _ = model.Update(tea.ColorProfileMsg{Profile: colorprofile.TrueColor}) + model.composer.SetValue("hi") + + _ = model.submit() + if model.welcomeAnimating { + t.Fatal("early prompt did not settle the welcome banner") + } +} + +func TestQuitLeavesNoPinnedFrame(t *testing.T) { + model := NewTUIModel(context.Background(), &fakeStreamer{}, TUIOptions{NoColor: true}) + model.width = 100 + model.height = 30 + model.Init() + model.quitCommand() + + if !model.quitting { + t.Fatal("quit command did not mark the model as quitting") + } + if view := model.View().Content; view != "" { + t.Errorf("quitting view = %q, want empty view", view) + } +} + +func TestWelcomeBannerWithNilUserStillShowsTitleAndHost(t *testing.T) { + model := NewTUIModel(context.Background(), &fakeStreamer{}, TUIOptions{ + NoColor: true, + APIHost: "api.rootly.com", + }) + + banner := model.welcomeBanner() + if !strings.Contains(banner, "Rootly AI") || !strings.Contains(banner, "api.rootly.com") { + t.Errorf("welcome banner = %q, want title and API host", banner) + } + if strings.Contains(banner, " · ") { + t.Errorf("welcome banner = %q, want no user row text when user is nil", banner) + } +} + +func TestTUIViewUsesAltScreenAndFixedFooterDuringConfirmation(t *testing.T) { + model := NewTUIModel(context.Background(), &fakeStreamer{}, TUIOptions{NoColor: true}) + view := model.View() + if !view.AltScreen || view.MouseMode != tea.MouseModeNone { + t.Errorf("view alt screen/mouse mode = %t/%v, want alt screen and no mouse capture", view.AltScreen, view.MouseMode) + } + model.Update(tea.WindowSizeMsg{Width: 80, Height: 24}) + model.state = tuiConfirming + model.pending = &rootly.AiChatStreamConfirmationRequired{Question: "Approve this action?"} + + content := model.View().Content + if got := lipgloss.Height(content); got != model.height { + t.Fatalf("confirming view height = %d, want %d", got, model.height) + } + if got := lastLine(content); !strings.Contains(got, "/ for commands · exit to quit") { + t.Errorf("last confirming row = %q, want fixed composer hint", got) + } + if !strings.Contains(content, "Approve this action?") { + t.Errorf("confirming view = %q, want confirmation card", content) + } +} + +func TestWelcomeHeaderContextUsesSessionBoundaryAndBinding(t *testing.T) { + sessionID := testSessionIDValue() + tests := []struct { + name string + options TUIOptions + want string + }{ + {name: "new", want: "New session"}, + {name: "resumed title", options: TUIOptions{SessionID: &sessionID, SessionTitle: "Incident review"}, want: "Incident review"}, + {name: "incident binding", options: TUIOptions{Incident: &api.Incident{SequentialID: "INC-123", Title: "Database latency"}}, want: "INC-123 · Database latency"}, + } + for _, test := range tests { + t.Run(test.name, func(t *testing.T) { + options := test.options + options.NoColor = true + model := NewTUIModel(context.Background(), &fakeStreamer{}, options) + lines := strings.Split(model.welcomeTextRows(), "\n") + if len(lines) < 4 { + t.Fatalf("header text rows = %d, want at least four: %q", len(lines), lines) + } + if last := lines[len(lines)-1]; last != test.want { + t.Errorf("header context = %q, want %q", last, test.want) + } + }) + } +} + +func TestWelcomeHeaderOmitsLogoBelowFiftyColumns(t *testing.T) { + model := NewTUIModel(context.Background(), &fakeStreamer{}, TUIOptions{NoColor: true}) + model.width = 49 + lines := strings.Split(strings.TrimSuffix(model.welcomeBanner(), "\n"), "\n") + if len(lines) != len(strings.Split(model.welcomeTextRows(), "\n")) { + t.Fatalf("narrow header rows = %q, want only the text rows", lines) + } + if strings.Contains(strings.Join(lines, "\n"), "▄") { + t.Errorf("narrow header includes logo: %q", lines) + } +} + +func TestTUIStatusVerbFooterAndSpinner(t *testing.T) { + now := time.Date(2025, time.May, 3, 15, 4, 5, 0, time.UTC) + model := NewTUIModel(context.Background(), &fakeStreamer{}, TUIOptions{NoColor: true}) + model.Update(tea.WindowSizeMsg{Width: 80, Height: 24}) + model.now = func() time.Time { return now } + model.verbIndex = func(int) int { return 3 } + sessionID := testSessionIDValue() + model.sessionID = &sessionID + model.startNewTurn() + now = now.Add(2 * time.Second) + + if got, want := model.turnFooter(), "✻ Untangled for 2s · done 3:04 PM"; got != want { + t.Errorf("turn footer = %q, want %q", got, want) + } + + model.state = tuiStreaming + model.startedAt = model.turnStartedAt + if view := ansi.Strip(model.View().Content); !strings.Contains(view, "Untangling…") { + t.Errorf("streaming view = %q, want active verb", view) + } + model.status = "Get incidents..." + view := ansi.Strip(model.View().Content) + if !strings.Contains(view, "Get incidents...") || strings.Contains(view, "Untangling…") { + t.Errorf("streaming view = %q, want server status to override active verb", view) + } + + now = model.turnStartedAt.Add(65 * time.Second) + if got, want := model.turnFooter(), "✻ Untangled for 1m 5s · done 3:05 PM"; got != want { + t.Errorf("long turn footer = %q, want %q", got, want) + } +} + +func currentUserForWelcome(name, email, fullNameWithTeam string) *rootly.CurrentUserResponse { + user := &rootly.CurrentUserResponse{} + user.Data.Attributes.Email = email + user.Data.Attributes.FullName = nullable.NewNullableWithValue(name) + if fullNameWithTeam != "" { + user.Data.Attributes.FullNameWithTeam = nullable.NewNullableWithValue(fullNameWithTeam) + } + return user +} + +func testSessionIDValue() uuid.UUID { + return uuid.MustParse(testSessionID) +} + +func lastLine(text string) string { + lines := strings.Split(text, "\n") + return lines[len(lines)-1] +} diff --git a/internal/api/ai_chat.go b/internal/api/ai_chat.go new file mode 100644 index 0000000..c17f9f3 --- /dev/null +++ b/internal/api/ai_chat.go @@ -0,0 +1,44 @@ +package api + +import ( + "context" + "net/url" + + openapi_types "github.com/oapi-codegen/runtime/types" + rootly "github.com/rootlyhq/rootly-go" +) + +func (c *Client) GetCurrentUserWithResponse(ctx context.Context) (*rootly.GetCurrentUserResponse, error) { + return c.client.GetCurrentUserWithResponse(ctx) +} + +func (c *Client) ListAiChatSessionsWithResponse(ctx context.Context, params *rootly.ListAiChatSessionsParams) (*rootly.ListAiChatSessionsResponse, error) { + return c.client.ListAiChatSessionsWithResponse(ctx, params) +} + +func (c *Client) ListAISkillsWithResponse(ctx context.Context, params *rootly.ListAISkillsParams, requestEditors ...rootly.RequestEditorFn) (*rootly.ListAISkillsResponse, error) { + return c.client.ListAISkillsWithResponse(ctx, params, requestEditors...) +} + +func (c *Client) ListAiChatSessionMessagesWithResponse(ctx context.Context, sessionID openapi_types.UUID, params *rootly.ListAiChatSessionMessagesParams) (*rootly.ListAiChatSessionMessagesResponse, error) { + return c.client.ListAiChatSessionMessagesWithResponse(ctx, sessionID, params) +} + +func (c *Client) UpdateAiChatSessionWithApplicationVndAPIPlusJSONBodyWithResponse(ctx context.Context, sessionID rootly.ID, body rootly.UpdateAiChatSessionApplicationVndAPIPlusJSONRequestBody) (*rootly.UpdateAiChatSessionResponse, error) { + return c.client.UpdateAiChatSessionWithApplicationVndAPIPlusJSONBodyWithResponse(ctx, sessionID, body) +} + +func (c *Client) DeleteAiChatSessionWithResponse(ctx context.Context, sessionID rootly.ID) (*rootly.DeleteAiChatSessionResponse, error) { + return c.client.DeleteAiChatSessionWithResponse(ctx, sessionID) +} + +func (c *Client) APIHost() string { + parsed, err := url.Parse(c.endpoint) + if err != nil { + return c.endpoint + } + if parsed.Host != "" { + return parsed.Host + } + return c.endpoint +} diff --git a/internal/api/ai_chat_stream_test.go b/internal/api/ai_chat_stream_test.go new file mode 100644 index 0000000..c8fc238 --- /dev/null +++ b/internal/api/ai_chat_stream_test.go @@ -0,0 +1,141 @@ +package api + +import ( + "context" + "fmt" + "iter" + "net/http" + "net/http/httptest" + "testing" + "time" + + "github.com/google/uuid" + rootly "github.com/rootlyhq/rootly-go" + + "github.com/rootlyhq/rootly-cli/internal/config" +) + +func TestStreamAiChatEvents(t *testing.T) { + sessionID := uuid.MustParse("39fd7596-152c-43d3-afce-48b792ae386d") + server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + if r.URL.Path != "/v1/ai/chat/stream" { + t.Errorf("path = %q, want /v1/ai/chat/stream", r.URL.Path) + } + if got := r.URL.Query().Get("message"); got != "Explain this" { + t.Errorf("message = %q, want Explain this", got) + } + if got := r.URL.Query().Get("session_id"); got != sessionID.String() { + t.Errorf("session_id = %q, want %q", got, sessionID) + } + if got := r.Header.Get("Accept"); got != "text/event-stream" { + t.Errorf("Accept = %q, want text/event-stream", got) + } + if got := r.Header.Get("Authorization"); got != "Bearer test-token" { + t.Errorf("Authorization = %q, want bearer token", got) + } + w.Header().Set("Content-Type", "text/event-stream") + _, _ = fmt.Fprint(w, "event: text\ndata: {\"content\":\"Hello\"}\n\n") + })) + defer server.Close() + + client, err := NewClient(&config.Config{APIKey: "test-token", Endpoint: server.URL}) + if err != nil { + t.Fatalf("NewClient() error = %v", err) + } + + message := "Explain this" + events, err := client.StreamAiChatEvents(context.Background(), &rootly.StreamAiChatParams{ + Message: &message, + SessionID: &sessionID, + }) + if err != nil { + t.Fatalf("StreamAiChatEvents() error = %v", err) + } + + count := 0 + for event, err := range events { + if err != nil { + t.Fatalf("stream event error = %v", err) + } + if event.Type != rootly.AiChatStreamEventText { + t.Errorf("event type = %q, want text", event.Type) + } + var payload rootly.AiChatStreamText + if err := event.Decode(&payload); err != nil { + t.Fatalf("Decode() error = %v", err) + } + if payload.Content != "Hello" { + t.Errorf("content = %q, want Hello", payload.Content) + } + count++ + } + if count != 1 { + t.Errorf("event count = %d, want 1", count) + } +} + +func TestDebugTransportDoesNotConsumeEventStream(t *testing.T) { + unblock := make(chan struct{}) + server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) { + w.Header().Set("Content-Type", "text/event-stream") + _, _ = fmt.Fprint(w, "event: text\ndata: {\"content\":\"first\"}\n\n") + w.(http.Flusher).Flush() + <-unblock + })) + defer server.Close() + defer close(unblock) + + client, err := NewClient(&config.Config{APIKey: "test-token", Endpoint: server.URL, Debug: true}) + if err != nil { + t.Fatalf("NewClient() error = %v", err) + } + message := "stream" + + type streamResult struct { + events iter.Seq2[rootly.AiChatStreamEvent, error] + err error + } + resultCh := make(chan streamResult, 1) + go func() { + events, err := client.StreamAiChatEvents(context.Background(), &rootly.StreamAiChatParams{Message: &message}) + resultCh <- streamResult{events: events, err: err} + }() + + var result streamResult + select { + case result = <-resultCh: + case <-time.After(time.Second): + t.Fatal("StreamAiChatEvents blocked while receiving response headers") + } + if result.err != nil { + t.Fatalf("StreamAiChatEvents() error = %v", result.err) + } + + eventCh := make(chan rootly.AiChatStreamEvent, 1) + errCh := make(chan error, 1) + go func() { + for event, err := range result.events { + if err != nil { + errCh <- err + return + } + eventCh <- event + return + } + }() + + select { + case event := <-eventCh: + var payload rootly.AiChatStreamText + if err := event.Decode(&payload); err != nil { + t.Fatalf("Decode() error = %v", err) + } + if payload.Content != "first" { + t.Errorf("content = %q, want first", payload.Content) + } + case err := <-errCh: + t.Fatalf("stream event error = %v", err) + case <-time.After(time.Second): + t.Fatal("first event was not readable before the server was unblocked") + } +} diff --git a/internal/api/ai_mcp.go b/internal/api/ai_mcp.go new file mode 100644 index 0000000..a5bc9ac --- /dev/null +++ b/internal/api/ai_mcp.go @@ -0,0 +1,194 @@ +package api + +import ( + "bytes" + "context" + "encoding/json" + "fmt" + "io" + "iter" + "net/http" + "strings" + + rootly "github.com/rootlyhq/rootly-go" + + "github.com/rootlyhq/rootly-cli/internal/mcp" +) + +type clientToolsKey struct{} + +// WithClientTools makes the next StreamAiChatEvents call send these local +// MCP tools as the `client_tools` JSON body. +func WithClientTools(ctx context.Context, tools []mcp.ClientTool) context.Context { + if len(tools) == 0 { + return ctx + } + return context.WithValue(ctx, clientToolsKey{}, tools) +} + +func clientToolsEditor(ctx context.Context) (rootly.RequestEditorFn, bool) { + tools, ok := ctx.Value(clientToolsKey{}).([]mcp.ClientTool) + if !ok || len(tools) == 0 { + return nil, false + } + return func(_ context.Context, req *http.Request) error { + body, err := json.Marshal(map[string]any{"client_tools": tools}) + if err != nil { + return err + } + req.Body = io.NopCloser(bytes.NewReader(body)) + req.GetBody = func() (io.ReadCloser, error) { return io.NopCloser(bytes.NewReader(body)), nil } + req.ContentLength = int64(len(body)) + req.Header.Set("Content-Type", "application/vnd.api+json") + return nil + }, true +} + +func (c *Client) streamAiChatEvents(ctx context.Context, params *rootly.StreamAiChatParams) (iter.Seq2[rootly.AiChatStreamEvent, error], error) { + if editor, ok := clientToolsEditor(ctx); ok { + return c.raw.StreamAiChatEvents(ctx, params, editor) + } + return c.raw.StreamAiChatEvents(ctx, params) +} + +type ToolResult struct { + SessionID string `json:"session_id"` + CallID string `json:"call_id"` + Output string `json:"output"` + IsError bool `json:"is_error"` +} + +// PostAiChatToolResult delivers a local tool's result to the blocked turn. +func (c *Client) PostAiChatToolResult(ctx context.Context, result ToolResult) error { + body, err := json.Marshal(result) + if err != nil { + return err + } + response, err := c.doJSON(ctx, http.MethodPost, "/v1/ai/chat/tool_results", body) + if err != nil { + return fmt.Errorf("send local tool result: %w", err) + } + defer func() { _ = response.Body.Close() }() + if response.StatusCode >= 200 && response.StatusCode < 300 { + _, _ = io.Copy(io.Discard, response.Body) + return nil + } + return fmt.Errorf("send local tool result: %s", apiErrorDetail(response)) +} + +type AIConnector struct { + ID string + Name string + ProviderType string + Status string + AvailableToAgent bool + ToolCount *int + ManageURL string +} + +type AIConnectorPage struct { + Connectors []AIConnector + NextPage int + StatusCode int +} + +// ListAIConnectors returns one page of the team's cloud MCP connectors. +func (c *Client) ListAIConnectors(ctx context.Context, page, pageSize int) (*AIConnectorPage, error) { + if page < 1 { + page = 1 + } + if pageSize < 1 || pageSize > 100 { + pageSize = 25 + } + response, err := c.doJSON(ctx, http.MethodGet, fmt.Sprintf("/v1/ai/connectors?page[number]=%d&page[size]=%d", page, pageSize), nil) + if err != nil { + return nil, fmt.Errorf("list connectors: %w", err) + } + defer func() { _ = response.Body.Close() }() + if response.StatusCode < 200 || response.StatusCode > 299 { + return &AIConnectorPage{StatusCode: response.StatusCode}, fmt.Errorf("list connectors: %s", apiErrorDetail(response)) + } + var document struct { + Data []struct { + ID string `json:"id"` + Attributes struct { + Name string `json:"name"` + ProviderType string `json:"provider_type"` + Status string `json:"status"` + AvailableToAgent bool `json:"available_to_agent"` + ToolCount *int `json:"tool_count"` + ManageURL string `json:"manage_url"` + } `json:"attributes"` + } `json:"data"` + Meta struct { + NextPage *int `json:"next_page"` + } `json:"meta"` + Links struct { + Next *string `json:"next"` + } `json:"links"` + } + if err := json.NewDecoder(io.LimitReader(response.Body, 8<<20)).Decode(&document); err != nil { + return nil, fmt.Errorf("list connectors: decode response: %w", err) + } + result := &AIConnectorPage{StatusCode: response.StatusCode} + for _, item := range document.Data { + attributes := item.Attributes + result.Connectors = append(result.Connectors, AIConnector{ + ID: item.ID, Name: attributes.Name, ProviderType: attributes.ProviderType, Status: attributes.Status, + AvailableToAgent: attributes.AvailableToAgent, ToolCount: attributes.ToolCount, ManageURL: attributes.ManageURL, + }) + } + switch { + case document.Meta.NextPage != nil && *document.Meta.NextPage > page: + result.NextPage = *document.Meta.NextPage + case document.Links.Next != nil && *document.Links.Next != "": + result.NextPage = page + 1 + case document.Meta.NextPage == nil && document.Links.Next == nil && len(document.Data) == pageSize: + result.NextPage = page + 1 + } + return result, nil +} + +func (c *Client) doJSON(ctx context.Context, method, path string, body []byte) (*http.Response, error) { + var reader io.Reader = http.NoBody + if body != nil { + reader = bytes.NewReader(body) + } + request, err := http.NewRequestWithContext(ctx, method, strings.TrimRight(c.endpoint, "/")+path, reader) + if err != nil { + return nil, err + } + if c.apiKey != "" { + request.Header.Set("Authorization", "Bearer "+c.apiKey) + } + request.Header.Set("Accept", "application/vnd.api+json, application/json") + if body != nil { + request.Header.Set("Content-Type", "application/vnd.api+json") + } + return c.httpClient.Do(request) +} + +func apiErrorDetail(response *http.Response) string { + data, _ := io.ReadAll(io.LimitReader(response.Body, 64*1024)) + var document struct { + Errors []struct { + Title string `json:"title"` + Detail string `json:"detail"` + } `json:"errors"` + Error string `json:"error"` + } + if json.Unmarshal(data, &document) == nil { + if len(document.Errors) > 0 { + if document.Errors[0].Detail != "" { + return fmt.Sprintf("HTTP %d: %s", response.StatusCode, document.Errors[0].Detail) + } + if document.Errors[0].Title != "" { + return fmt.Sprintf("HTTP %d: %s", response.StatusCode, document.Errors[0].Title) + } + } + if document.Error != "" { + return fmt.Sprintf("HTTP %d: %s", response.StatusCode, document.Error) + } + } + return fmt.Sprintf("HTTP %d", response.StatusCode) +} diff --git a/internal/api/ai_mcp_test.go b/internal/api/ai_mcp_test.go new file mode 100644 index 0000000..18bbd4a --- /dev/null +++ b/internal/api/ai_mcp_test.go @@ -0,0 +1,92 @@ +package api + +import ( + "context" + "encoding/json" + "io" + "net/http" + "net/http/httptest" + "strings" + "testing" + + rootly "github.com/rootlyhq/rootly-go" + + "github.com/rootlyhq/rootly-cli/internal/config" + "github.com/rootlyhq/rootly-cli/internal/mcp" +) + +func TestStreamAddsClientToolsBody(t *testing.T) { + var body map[string]any + var query, contentType string + server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + contentType = r.Header.Get("Content-Type") + data, _ := io.ReadAll(r.Body) + _ = json.Unmarshal(data, &body) + query = r.URL.Query().Get("message") + w.Header().Set("Content-Type", "text/event-stream") + _, _ = io.WriteString(w, "event: done\ndata: {}\n\n") + })) + defer server.Close() + client, err := NewClient(&config.Config{APIKey: "test-token", Endpoint: server.URL}) + if err != nil { + t.Fatal(err) + } + message := "hi" + tools := []mcp.ClientTool{{Name: "local_fs_read", Description: "read", InputSchema: json.RawMessage(`{"type":"object"}`), ReadOnly: true}} + events, err := client.StreamAiChatEvents(WithClientTools(context.Background(), tools), &rootly.StreamAiChatParams{Message: &message}) + if err != nil { + t.Fatal(err) + } + for range events { + } + sent, _ := body["client_tools"].([]any) + if query != "hi" || len(sent) != 1 || sent[0].(map[string]any)["read_only"] != true || contentType != "application/vnd.api+json" { + t.Fatalf("query=%q content-type=%q body=%v", query, contentType, body) + } +} + +func TestPostToolResultAndConnectors(t *testing.T) { + var posted ToolResult + var connectorQuery, postedContentType string + server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + switch r.URL.Path { + case "/v1/ai/chat/tool_results": + postedContentType = r.Header.Get("Content-Type") + _ = json.NewDecoder(r.Body).Decode(&posted) + if posted.CallID == "missing" { + w.Header().Set("Content-Type", "application/vnd.api+json") + w.WriteHeader(http.StatusNotFound) + _, _ = io.WriteString(w, `{"errors":[{"title":"Unknown call"}]}`) + return + } + w.WriteHeader(http.StatusAccepted) + case "/v1/ai/connectors": + connectorQuery = r.URL.RawQuery + _, _ = io.WriteString(w, `{"data":[{"id":"c1","type":"ai_connectors","attributes":{"name":"Datadog","provider_type":"datadog","status":"connected","available_to_agent":true,"tool_count":12,"manage_url":"https://rootly.com/x"}},{"id":"c2","type":"ai_connectors","attributes":{"name":"Custom","provider_type":"custom","status":"error","available_to_agent":false,"tool_count":null,"manage_url":""}}],"meta":{"current_page":1,"next_page":2,"total_pages":2}}`) + } + })) + defer server.Close() + client, err := NewClient(&config.Config{APIKey: "test-token", Endpoint: server.URL}) + if err != nil { + t.Fatal(err) + } + if err := client.PostAiChatToolResult(context.Background(), ToolResult{SessionID: "s", CallID: "c", Output: "User declined", IsError: true}); err != nil { + t.Fatal(err) + } + if posted.Output != "User declined" || !posted.IsError || posted.SessionID != "s" || postedContentType != "application/vnd.api+json" { + t.Fatalf("posted = %+v", posted) + } + if err := client.PostAiChatToolResult(context.Background(), ToolResult{CallID: "missing"}); err == nil || !strings.Contains(err.Error(), "Unknown call") { + t.Fatalf("err = %v", err) + } + page, err := client.ListAIConnectors(context.Background(), 1, 25) + if err != nil { + t.Fatal(err) + } + if !strings.Contains(connectorQuery, "page%5Bnumber%5D=1") && !strings.Contains(connectorQuery, "page[number]=1") { + t.Fatalf("query = %q", connectorQuery) + } + if len(page.Connectors) != 2 || page.NextPage != 2 || *page.Connectors[0].ToolCount != 12 || page.Connectors[1].ToolCount != nil { + t.Fatalf("page = %+v", page) + } +} diff --git a/internal/api/client.go b/internal/api/client.go index e8764e5..a1fa2e2 100644 --- a/internal/api/client.go +++ b/internal/api/client.go @@ -5,6 +5,7 @@ import ( "encoding/json" "fmt" "io" + "iter" "net/http" "net/http/httputil" neturl "net/url" @@ -24,6 +25,7 @@ var Version = "dev" type Client struct { client *rootly.ClientWithResponses + raw *rootly.Client endpoint string apiKey string httpClient *http.Client @@ -39,6 +41,16 @@ func (dt *debugTransport) RoundTrip(req *http.Request) (*http.Response, error) { if auth := safeReq.Header.Get("Authorization"); auth != "" { safeReq.Header.Set("Authorization", "Bearer [REDACTED]") } + if safeReq.URL != nil { + query := safeReq.URL.Query() + if values, ok := query["message"]; ok { + for index := range values { + values[index] = "[REDACTED]" + } + query["message"] = values + safeReq.URL.RawQuery = query.Encode() + } + } safeReq.Body = http.NoBody dump, _ := httputil.DumpRequestOut(safeReq, false) fmt.Fprintf(os.Stderr, "\n--- DEBUG REQUEST ---\n%s\n", dump) @@ -49,7 +61,8 @@ func (dt *debugTransport) RoundTrip(req *http.Request) (*http.Response, error) { return resp, err } - dump, _ = httputil.DumpResponse(resp, true) + includeBody := !strings.HasPrefix(strings.ToLower(resp.Header.Get("Content-Type")), "text/event-stream") + dump, _ = httputil.DumpResponse(resp, includeBody) fmt.Fprintf(os.Stderr, "--- DEBUG RESPONSE ---\n%s\n", dump) return resp, err @@ -925,22 +938,28 @@ func NewClient(cfg *config.Config) (*Client, error) { } } - client, err := rootly.NewClientWithResponses(endpoint, + raw, err := rootly.NewClient(endpoint, rootly.WithHTTPClient(httpClient), rootly.WithRequestEditorFn(reqEditorFn), ) if err != nil { return nil, fmt.Errorf("failed to create rootly client: %w", err) } + client := &rootly.ClientWithResponses{ClientInterface: raw} return &Client{ client: client, + raw: raw, endpoint: endpoint, apiKey: cfg.APIKey, httpClient: httpClient, }, nil } +func (c *Client) StreamAiChatEvents(ctx context.Context, params *rootly.StreamAiChatParams) (iter.Seq2[rootly.AiChatStreamEvent, error], error) { + return c.streamAiChatEvents(ctx, params) +} + // ensureScheme adds a scheme if missing, using http for localhost/127.0.0.1. // For localhost without an explicit path, it also appends /api since the // Rails monolith serves the API under /api/v1 rather than /v1. diff --git a/internal/api/client_test.go b/internal/api/client_test.go index db50d74..8d108c3 100644 --- a/internal/api/client_test.go +++ b/internal/api/client_test.go @@ -267,6 +267,58 @@ func TestDebugTransportRedactsAuthorization(t *testing.T) { } } +func TestDebugTransportRedactsMessageQueryParameter(t *testing.T) { + server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + if values := r.URL.Query()["message"]; len(values) == 2 && values[0] == "first-private-prompt" && values[1] == "second-private-prompt" && r.URL.Query().Get("other") == "visible" { + w.Header().Set("X-Original-Query-Preserved", "true") + } + w.WriteHeader(http.StatusOK) + })) + defer server.Close() + + dt := &debugTransport{transport: http.DefaultTransport} + req, err := http.NewRequestWithContext(context.Background(), "GET", server.URL+"/v1/ai/chat/stream?message=first-private-prompt&other=visible&message=second-private-prompt", http.NoBody) + if err != nil { + t.Fatal(err) + } + + oldStderr := os.Stderr + reader, writer, err := os.Pipe() + if err != nil { + t.Fatal(err) + } + os.Stderr = writer + resp, roundTripErr := dt.RoundTrip(req) + writer.Close() + os.Stderr = oldStderr + if roundTripErr != nil { + reader.Close() + t.Fatalf("RoundTrip failed: %v", roundTripErr) + } + resp.Body.Close() + outputBytes, err := io.ReadAll(reader) + reader.Close() + if err != nil { + t.Fatal(err) + } + + output := string(outputBytes) + for _, privateValue := range []string{"first-private-prompt", "second-private-prompt"} { + if strings.Contains(output, privateValue) { + t.Errorf("debug output contains private message query value %q: %s", privateValue, output) + } + } + if got := strings.Count(output, "message=%5BREDACTED%5D"); got != 2 { + t.Errorf("redacted message query count = %d, want two: %s", got, output) + } + if !strings.Contains(output, "other=visible") { + t.Errorf("debug output = %s, want unrelated query parameter preserved", output) + } + if got := resp.Header.Get("X-Original-Query-Preserved"); got != "true" { + t.Errorf("original query preserved = %q, want true", got) + } +} + func TestDebugTransportPreservesActualAuthHeader(t *testing.T) { var receivedAuth string server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { diff --git a/internal/cmd/agent/agent.go b/internal/cmd/agent/agent.go new file mode 100644 index 0000000..e85c018 --- /dev/null +++ b/internal/cmd/agent/agent.go @@ -0,0 +1,332 @@ +package agent + +import ( + "context" + "errors" + "fmt" + "io" + "os" + "time" + + tea "charm.land/bubbletea/v2" + "github.com/charmbracelet/colorprofile" + "github.com/charmbracelet/x/ansi" + "github.com/google/uuid" + "github.com/mattn/go-isatty" + rootly "github.com/rootlyhq/rootly-go" + "github.com/spf13/cobra" + "github.com/spf13/viper" + + aiagent "github.com/rootlyhq/rootly-cli/internal/agent" + "github.com/rootlyhq/rootly-cli/internal/api" + "github.com/rootlyhq/rootly-cli/internal/config" + "github.com/rootlyhq/rootly-cli/internal/mcp" + "github.com/rootlyhq/rootly-cli/internal/oauth" +) + +var AgentCmd = &cobra.Command{ + Use: "agent", + Aliases: []string{"chat"}, + Short: "Chat with Rootly AI", + Long: `Chat with Rootly AI. + +Print mode (-p, or piped stdin) exit codes: + 0 success + 1 error + 2 not authenticated (run 'rootly login' or set ROOTLY_API_KEY)`, + Hidden: true, + Args: cobra.NoArgs, + RunE: func(cmd *cobra.Command, _ []string) error { + err := runAgent(cmd) + if isPrintMode(cmd) { + return aiagent.PrintExitError(err) + } + return err + }, +} + +func isPrintMode(cmd *cobra.Command) bool { + return cmd.Flags().Changed("print") || !isTerminalReader(cmd.InOrStdin()) +} + +func runAgent(cmd *cobra.Command) error { + prompt, _ := cmd.Flags().GetString("print") + stdin := cmd.InOrStdin() + resumeValue, _ := cmd.Flags().GetString("resume") + continueSession, _ := cmd.Flags().GetBool("continue") + incidentValue, _ := cmd.Flags().GetString("incident") + alertValue, _ := cmd.Flags().GetString("alert") + outputFormat, _ := cmd.Flags().GetString("output-format") + printMode := isPrintMode(cmd) + if cmd.Flags().Changed("output-format") && !printMode { + return fmt.Errorf("--output-format is only valid with print or stdin mode") + } + if outputFormat != "text" && outputFormat != "json" && outputFormat != "stream-json" { + return fmt.Errorf("invalid output format %q (expected text, json, or stream-json)", outputFormat) + } + var ( + uiWriter io.Writer + err error + ) + if !printMode { + uiWriter, err = agentUIWriter(cmd.OutOrStdout(), cmd.ErrOrStderr(), isTerminalWriter) + if err != nil { + return err + } + } + client, err := getAPIClient() + if err != nil { + return err + } + var sessionID *uuid.UUID + sessionTitle := "" + if resumeValue != "" { + id, parseErr := uuid.Parse(resumeValue) + if parseErr != nil { + return fmt.Errorf("invalid resume session UUID %q: %w", resumeValue, parseErr) + } + sessionID = &id + } + currentUser, err := preflightAI(cmd, client) + if err != nil { + return err + } + if continueSession { + id, title, err := latestSession(cmd.Context(), client) + if err != nil { + return aiagent.SanitizeTerminalError(err) + } + sessionID = &id + sessionTitle = title + } + var incident *api.Incident + if incidentValue != "" { + incident, err = client.GetIncidentByID(cmd.Context(), api.NormalizeIncidentID(incidentValue)) + if err != nil { + return fmt.Errorf("incident not found: %w", aiagent.SanitizeTerminalError(err)) + } + } + var boundAlert *api.Alert + if alertValue != "" { + boundAlert, err = client.GetAlertByID(cmd.Context(), alertValue) + if err != nil { + return fmt.Errorf("alert not found: %w", aiagent.SanitizeTerminalError(err)) + } + } + manager := startLocalMCP(cmd) + defer manager.Close() + if printMode { + waitForLocalMCP(cmd.Context(), manager, 20*time.Second) + return aiagent.RunPrint(cmd.Context(), client, aiagent.PrintOptions{ + Prompt: prompt, + Stdin: stdin, + Stdout: cmd.OutOrStdout(), + Stderr: cmd.ErrOrStderr(), + NoColor: viper.GetBool("no-color"), + Quiet: viper.GetBool("quiet"), + APIHost: client.APIHost(), + SessionID: sessionID, + OutputFormat: outputFormat, + Incident: incident, + Alert: boundAlert, + IsStdinTTY: func() bool { return isTerminalReader(stdin) }, + IsStdoutTTY: func() bool { return isTerminalWriter(cmd.OutOrStdout()) }, + IsStderrTTY: func() bool { return isTerminalWriter(cmd.ErrOrStderr()) }, + LocalTools: manager, + ToolResults: client, + }) + } + + noColor := viper.GetBool("no-color") + model := aiagent.NewTUIModel(cmd.Context(), client, aiagent.TUIOptions{ + NoColor: noColor, Version: cmd.Root().Version, API: client, SessionID: sessionID, SessionTitle: sessionTitle, + APIHost: client.APIHost(), CurrentUser: currentUser, WorkingDir: agentWorkingDir(), + Incident: incident, Alert: boundAlert, + MCP: manager, ToolResults: client, + }) + options := []tea.ProgramOption{ + tea.WithInput(stdin), + tea.WithOutput(uiWriter), + tea.WithContext(cmd.Context()), + } + if noColor { + options = append(options, tea.WithColorProfile(colorprofile.Ascii)) + } + noClear, _ := cmd.Flags().GetBool("no-clear") + if err := clearTerminalForAgent(uiWriter, !noClear, isTerminalWriter(uiWriter)); err != nil { + return err + } + program := tea.NewProgram(model, options...) + return runAgentTUI(cmd.OutOrStdout(), uiWriter, program.Run) +} + +func init() { + AgentCmd.Flags().StringP("print", "p", "", "Send a prompt and print the response") + AgentCmd.Flags().StringP("resume", "r", "", "Resume an AI chat session by UUID") + AgentCmd.Flags().BoolP("continue", "c", false, "Continue the most recent AI chat session (incident/alert chats: use --resume )") + AgentCmd.Flags().String("incident", "", "Bind the chat to an incident") + AgentCmd.Flags().String("alert", "", "Bind the chat to an alert") + AgentCmd.Flags().String("output-format", "text", "Print output format: text, json, or stream-json") + AgentCmd.Flags().Bool("no-clear", false, "Don't clear the terminal and scrollback on start") + AgentCmd.MarkFlagsMutuallyExclusive("resume", "continue", "incident", "alert") +} + +func latestSession(ctx context.Context, client *api.Client) (uuid.UUID, string, error) { + limit, archived := 1, false + response, err := client.ListAiChatSessionsWithResponse(ctx, &rootly.ListAiChatSessionsParams{Limit: &limit, Archived: &archived}) + if err != nil { + return uuid.Nil, "", aiagent.SanitizeTerminalError(err) + } + if response == nil { + return uuid.Nil, "", errors.New("empty response while listing AI chat sessions") + } + if response.ApplicationVndAPIJSON200 == nil { + return uuid.Nil, "", aiagent.StatusResponseError(response.StatusCode(), response.Body) + } + sessions := response.ApplicationVndAPIJSON200.Sessions + if len(sessions) == 0 { + return uuid.Nil, "", errors.New("no previous AI chat session to continue (incident/alert-bound chats are resumed with --resume )") + } + id, err := uuid.Parse(sessions[0].ID.String()) + if err != nil { + return uuid.Nil, "", fmt.Errorf("invalid session UUID from API: %w", err) + } + return id, sessions[0].Title, nil +} + +func preflightAI(cmd *cobra.Command, client *api.Client) (*rootly.CurrentUserResponse, error) { + response, err := client.GetCurrentUserWithResponse(cmd.Context()) + if err != nil { + if viper.GetBool("debug") { + _, _ = fmt.Fprintf(cmd.ErrOrStderr(), "AI preflight failed: %v\n", aiagent.SanitizeTerminalError(err)) + } + return nil, nil + } + if response == nil { + if viper.GetBool("debug") { + _, _ = fmt.Fprintln(cmd.ErrOrStderr(), "AI preflight returned an empty response") + } + return nil, nil + } + if response.StatusCode() == 401 { + return nil, aiagent.StatusResponseError(response.StatusCode(), response.Body) + } + if response.StatusCode() != 200 || response.ApplicationVndAPIJSON200 == nil { + if response.StatusCode() == 403 { + mapped := aiagent.StatusResponseError(response.StatusCode(), response.Body) + if aiagent.IsAIChatDisabled(mapped) { + return nil, mapped + } + } + if viper.GetBool("debug") { + _, _ = fmt.Fprintf(cmd.ErrOrStderr(), "AI preflight returned HTTP %d\n", response.StatusCode()) + } + return nil, nil + } + current := response.ApplicationVndAPIJSON200 + if enabled, exists := current.Data.Attributes.FeatureFlags["ai-chat"]; exists && !enabled { + return nil, aiagent.AIChatDisabledError() + } + return current, nil +} + +func getAPIClient() (*api.Client, error) { + token := viper.GetString("api_key") + if token == "" && !oauth.HasTokens() { + return nil, aiagent.NotAuthenticatedError("authentication required: run 'rootly login' or set ROOTLY_API_KEY") + } + endpoint := viper.GetString("api_host") + if endpoint == "" { + endpoint = config.DefaultEndpoint + } + return api.NewClient(&config.Config{ + APIKey: token, + Endpoint: endpoint, + Debug: viper.GetBool("debug"), + }) +} + +func runAgentTUI(stdout, uiWriter io.Writer, run func() (tea.Model, error)) error { + finalModel, runErr := run() + if model, ok := finalModel.(*aiagent.TUIModel); ok { + transcript := model.Transcript() + stdoutTranscript := transcript + if !isTerminalWriter(stdout) { + stdoutTranscript = ansi.Strip(transcript) + } + if _, err := io.WriteString(stdout, stdoutTranscript); err != nil && runErr == nil { + runErr = err + } + if uiWriter != stdout { + if _, err := io.WriteString(uiWriter, transcript); err != nil && runErr == nil { + runErr = err + } + } + } + return runErr +} + +func isTerminalReader(reader io.Reader) bool { + file, ok := reader.(*os.File) + return ok && isatty.IsTerminal(file.Fd()) +} + +func isTerminalWriter(writer io.Writer) bool { + file, ok := writer.(*os.File) + return ok && isatty.IsTerminal(file.Fd()) +} + +func agentUIWriter(stdout, stderr io.Writer, isTTY func(io.Writer) bool) (io.Writer, error) { + if isTTY(stdout) { + return stdout, nil + } + if isTTY(stderr) { + return stderr, nil + } + return nil, errors.New(`interactive mode needs a terminal; use -p "" for non-interactive output`) +} + +func clearTerminalForAgent(writer io.Writer, enabled, terminal bool) error { + if !enabled || !terminal { + return nil + } + _, err := io.WriteString(writer, "\x1b[2J\x1b[3J\x1b[H") + return err +} + +func startLocalMCP(cmd *cobra.Command) *mcp.Manager { + cwd, err := os.Getwd() + if err != nil { + cwd = "." + } + manager := mcp.NewManager(mcp.ManagerOptions{CWD: cwd, Version: cmd.Root().Version}) + if err := manager.Load(); err != nil && !viper.GetBool("quiet") { + _, _ = fmt.Fprintln(cmd.ErrOrStderr(), aiagent.SanitizeTerminalError(fmt.Errorf("local MCP config: %w", err))) + } + manager.Start() + return manager +} + +func waitForLocalMCP(ctx context.Context, manager *mcp.Manager, timeout time.Duration) { + deadline := time.NewTimer(timeout) + defer deadline.Stop() + ticker := time.NewTicker(100 * time.Millisecond) + defer ticker.Stop() + for manager.Connecting() > 0 { + select { + case <-ctx.Done(): + return + case <-deadline.C: + return + case <-ticker.C: + } + } +} + +func agentWorkingDir() string { + dir, err := os.Getwd() + if err != nil { + return "" + } + return dir +} diff --git a/internal/cmd/agent/agent_test.go b/internal/cmd/agent/agent_test.go new file mode 100644 index 0000000..e8cf019 --- /dev/null +++ b/internal/cmd/agent/agent_test.go @@ -0,0 +1,461 @@ +package agent + +import ( + "bytes" + "context" + "encoding/json" + "errors" + "io" + "net/http" + "net/http/httptest" + "strings" + "sync/atomic" + "testing" + + tea "charm.land/bubbletea/v2" + "github.com/charmbracelet/x/ansi" + "github.com/spf13/cobra" + "github.com/spf13/viper" + + aiagent "github.com/rootlyhq/rootly-cli/internal/agent" + "github.com/rootlyhq/rootly-cli/internal/api" + "github.com/rootlyhq/rootly-cli/internal/config" +) + +func TestPreflightAIFeatureFlag(t *testing.T) { + tests := []struct { + name string + status int + flags map[string]bool + wantError string + wantUser bool + wantAIChatDisabled bool + }{ + {name: "disabled", status: http.StatusOK, flags: map[string]bool{"ai-chat": false}, wantError: "Rootly AI chat isn't enabled for your organization. Contact your Rootly account team or Rootly support to turn it on.", wantAIChatDisabled: true}, + {name: "enabled", status: http.StatusOK, flags: map[string]bool{"ai-chat": true}, wantUser: true}, + {name: "flag absent", status: http.StatusOK, flags: map[string]bool{"ai-native-home": true}, wantUser: true}, + {name: "unauthorized", status: http.StatusUnauthorized, wantError: "not authenticated: run `rootly login` or set ROOTLY_API_KEY"}, + {name: "other failure proceeds", status: http.StatusInternalServerError}, + {name: "disabled response", status: http.StatusForbidden, wantError: "Rootly AI chat isn't enabled for your organization. Contact your Rootly account team or Rootly support to turn it on.", wantAIChatDisabled: true}, + } + for _, test := range tests { + t.Run(test.name, func(t *testing.T) { + server := httptest.NewServer(http.HandlerFunc(func(writer http.ResponseWriter, request *http.Request) { + if request.URL.Path != "/v1/users/me" { + t.Errorf("preflight path = %q, want /v1/users/me", request.URL.Path) + } + writer.Header().Set("Content-Type", "application/vnd.api+json") + writer.WriteHeader(test.status) + switch test.status { + case http.StatusOK: + _ = json.NewEncoder(writer).Encode(map[string]any{ + "data": map[string]any{ + "id": "user-1", "type": "users", + "attributes": map[string]any{"email": "test@example.com", "feature_flags": test.flags}, + }, + }) + case http.StatusUnauthorized: + _, _ = writer.Write([]byte(`{"errors":[{"title":"Unauthorized"}]}`)) + case http.StatusForbidden: + _, _ = writer.Write([]byte(`{"errors":[{"title":"AI chat is not enabled for this team"}]}`)) + } + })) + defer server.Close() + + client, err := api.NewClient(&config.Config{APIKey: "test-key", Endpoint: server.URL}) + if err != nil { + t.Fatalf("NewClient() error = %v", err) + } + var stderr bytes.Buffer + cmd := &cobra.Command{} + cmd.SetContext(context.Background()) + cmd.SetErr(&stderr) + currentUser, err := preflightAI(cmd, client) + if test.wantError != "" { + if err == nil || err.Error() != test.wantError { + t.Fatalf("preflightAI() error = %v, want %q", err, test.wantError) + } + if got := aiagent.IsAIChatDisabled(err); got != test.wantAIChatDisabled { + t.Errorf("IsAIChatDisabled(error) = %v, want %v", got, test.wantAIChatDisabled) + } + if got, want := aiagent.IsNotAuthenticated(err), test.status == http.StatusUnauthorized; got != want { + t.Errorf("IsNotAuthenticated(error) = %v, want %v", got, want) + } + return + } + if err != nil { + t.Fatalf("preflightAI() error = %v", err) + } + if (currentUser != nil) != test.wantUser { + t.Errorf("current user returned = %v, want %v", currentUser != nil, test.wantUser) + } + }) + } +} + +func TestAgentPreflightRunsBeforeIncidentLookup(t *testing.T) { + var incidentLookups atomic.Int32 + server := httptest.NewServer(http.HandlerFunc(func(writer http.ResponseWriter, request *http.Request) { + if request.URL.Path == "/v1/users/me" { + writer.Header().Set("Content-Type", "application/vnd.api+json") + _, _ = writer.Write([]byte(`{"data":{"id":"user-1","type":"users","attributes":{"feature_flags":{"ai-chat":false}}}}`)) + return + } + if strings.Contains(request.URL.Path, "/incidents") { + incidentLookups.Add(1) + } + http.NotFound(writer, request) + })) + defer server.Close() + + oldAPIKey, oldAPIHost, oldDebug := viper.GetString("api_key"), viper.GetString("api_host"), viper.GetBool("debug") + viper.Set("api_key", "test-key") + viper.Set("api_host", server.URL) + viper.Set("debug", false) + t.Cleanup(func() { + viper.Set("api_key", oldAPIKey) + viper.Set("api_host", oldAPIHost) + viper.Set("debug", oldDebug) + }) + + cmd := &cobra.Command{} + cmd.Flags().StringP("print", "p", "", "") + cmd.Flags().StringP("resume", "r", "", "") + cmd.Flags().BoolP("continue", "c", false, "") + cmd.Flags().String("incident", "", "") + cmd.Flags().String("alert", "", "") + cmd.Flags().String("output-format", "text", "") + cmd.SetContext(context.Background()) + cmd.SetIn(bytes.NewReader(nil)) + cmd.SetOut(&bytes.Buffer{}) + cmd.SetErr(&bytes.Buffer{}) + if err := cmd.Flags().Set("incident", "INC-8"); err != nil { + t.Fatalf("set incident flag: %v", err) + } + + err := AgentCmd.RunE(cmd, nil) + wantError := "Rootly AI chat isn't enabled for your organization. Contact your Rootly account team or Rootly support to turn it on." + if err == nil || err.Error() != wantError { + t.Fatalf("AgentCmd.RunE() error = %v, want %q", err, wantError) + } + if got := incidentLookups.Load(); got != 0 { + t.Errorf("incident lookup requests = %d, want none before AI preflight", got) + } +} + +func TestLatestSessionRequestsMostRecentUnarchivedSession(t *testing.T) { + sessionID := "8e9bd763-7c86-4741-9766-86a27d3c36d1" + server := httptest.NewServer(http.HandlerFunc(func(writer http.ResponseWriter, request *http.Request) { + if request.URL.Path != "/v1/ai/chat/sessions" { + t.Errorf("list path = %q, want /v1/ai/chat/sessions", request.URL.Path) + } + if request.URL.Query().Get("limit") != "1" || request.URL.Query().Get("archived") != "false" { + t.Errorf("list query = %q, want limit=1 and archived=false", request.URL.RawQuery) + } + writer.Header().Set("Content-Type", "application/vnd.api+json") + _, _ = writer.Write([]byte(`{"meta":{"pagination":{"has_more":false,"next_before":null}},"sessions":[{"id":"` + sessionID + `","title":"Recent","last_active_at":"2026-10-29T12:00:00Z"}]}`)) + })) + defer server.Close() + + client, err := api.NewClient(&config.Config{APIKey: "test-key", Endpoint: server.URL}) + if err != nil { + t.Fatalf("NewClient() error = %v", err) + } + gotID, title, err := latestSession(context.Background(), client) + if err != nil { + t.Fatalf("latestSession() error = %v", err) + } + if gotID.String() != sessionID || title != "Recent" { + t.Errorf("latestSession() = %s, %q, want %s, Recent", gotID, title, sessionID) + } +} + +func TestLatestSessionErrors(t *testing.T) { + tests := []struct { + name string + status int + body string + wantError string + }{ + {name: "no sessions", status: http.StatusOK, body: `{"meta":{"pagination":{"has_more":false,"next_before":null}},"sessions":[]}`, wantError: "no previous AI chat session to continue (incident/alert-bound chats are resumed with --resume )"}, + {name: "disabled", status: http.StatusForbidden, body: `{"errors":[{"title":"AI chat is not enabled for this team"}]}`, wantError: "Rootly AI chat isn't enabled for your organization. Contact your Rootly account team or Rootly support to turn it on."}, + } + for _, test := range tests { + t.Run(test.name, func(t *testing.T) { + server := httptest.NewServer(http.HandlerFunc(func(writer http.ResponseWriter, _ *http.Request) { + writer.Header().Set("Content-Type", "application/vnd.api+json") + writer.WriteHeader(test.status) + _, _ = writer.Write([]byte(test.body)) + })) + defer server.Close() + client, err := api.NewClient(&config.Config{APIKey: "test-key", Endpoint: server.URL}) + if err != nil { + t.Fatalf("NewClient() error = %v", err) + } + _, _, err = latestSession(context.Background(), client) + if err == nil || err.Error() != test.wantError { + t.Errorf("latestSession() error = %v, want %q", err, test.wantError) + } + }) + } +} + +func TestContinueFlagHelp(t *testing.T) { + flag := AgentCmd.Flags().Lookup("continue") + if flag == nil { + t.Fatal("continue flag is not registered") + } + want := "Continue the most recent AI chat session (incident/alert chats: use --resume )" + if flag.Usage != want { + t.Errorf("continue help = %q, want %q", flag.Usage, want) + } +} + +func TestAgentUIWriter(t *testing.T) { + var stdout, stderr bytes.Buffer + tests := []struct { + name string + stdoutTTY bool + stderrTTY bool + wantWriter io.Writer + wantError string + }{ + {name: "stdout terminal", stdoutTTY: true, stderrTTY: true, wantWriter: &stdout}, + {name: "stderr terminal fallback", stderrTTY: true, wantWriter: &stderr}, + {name: "no terminal", wantError: `interactive mode needs a terminal; use -p "" for non-interactive output`}, + } + for _, test := range tests { + t.Run(test.name, func(t *testing.T) { + isTTY := func(writer io.Writer) bool { + switch writer { + case &stdout: + return test.stdoutTTY + case &stderr: + return test.stderrTTY + default: + return false + } + } + got, err := agentUIWriter(&stdout, &stderr, isTTY) + if test.wantError != "" { + if err == nil || err.Error() != test.wantError { + t.Fatalf("agentUIWriter() error = %v, want %q", err, test.wantError) + } + if got != nil { + t.Errorf("agentUIWriter() writer = %v, want nil", got) + } + return + } + if err != nil { + t.Fatalf("agentUIWriter() error = %v", err) + } + if got != test.wantWriter { + t.Errorf("agentUIWriter() writer = %v, want %v", got, test.wantWriter) + } + }) + } +} + +func TestClearTerminalForAgent(t *testing.T) { + tests := []struct { + name string + enabled bool + terminal bool + want string + }{ + {name: "terminal", enabled: true, terminal: true, want: "\x1b[2J\x1b[3J\x1b[H"}, + {name: "disabled", enabled: false, terminal: true}, + {name: "non-terminal", enabled: true}, + } + for _, test := range tests { + t.Run(test.name, func(t *testing.T) { + var output bytes.Buffer + if err := clearTerminalForAgent(&output, test.enabled, test.terminal); err != nil { + t.Fatalf("clearTerminalForAgent() error = %v", err) + } + if got := output.String(); got != test.want { + t.Errorf("clear output = %q, want %q", got, test.want) + } + }) + } +} + +func TestAgentPrintModeDoesNotEmitTerminalClear(t *testing.T) { + server := httptest.NewServer(http.HandlerFunc(func(writer http.ResponseWriter, request *http.Request) { + switch request.URL.Path { + case "/v1/users/me": + writer.Header().Set("Content-Type", "application/vnd.api+json") + _, _ = writer.Write([]byte(`{"data":{"id":"user-1","type":"users","attributes":{"email":"test@example.com","feature_flags":{"ai-chat":true}}}}`)) + case "/v1/ai/chat/stream": + writer.Header().Set("Content-Type", "text/event-stream") + _, _ = writer.Write([]byte("event: text\ndata: {\"content\":\"answer\"}\n\nevent: done\ndata: {\"status\":\"complete\"}\n\n")) + default: + http.NotFound(writer, request) + } + })) + defer server.Close() + + oldAPIKey, oldAPIHost := viper.GetString("api_key"), viper.GetString("api_host") + oldNoColor, oldQuiet := viper.GetBool("no-color"), viper.GetBool("quiet") + viper.Set("api_key", "test-key") + viper.Set("api_host", server.URL) + viper.Set("no-color", true) + viper.Set("quiet", false) + t.Cleanup(func() { + viper.Set("api_key", oldAPIKey) + viper.Set("api_host", oldAPIHost) + viper.Set("no-color", oldNoColor) + viper.Set("quiet", oldQuiet) + }) + + cmd := &cobra.Command{} + cmd.Flags().StringP("print", "p", "", "") + cmd.Flags().StringP("resume", "r", "", "") + cmd.Flags().BoolP("continue", "c", false, "") + cmd.Flags().String("incident", "", "") + cmd.Flags().String("alert", "", "") + cmd.Flags().String("output-format", "text", "") + if err := cmd.Flags().Set("print", "question"); err != nil { + t.Fatalf("set print flag: %v", err) + } + cmd.SetContext(context.Background()) + cmd.SetIn(bytes.NewBuffer(nil)) + var stdout, stderr bytes.Buffer + cmd.SetOut(&stdout) + cmd.SetErr(&stderr) + + if err := AgentCmd.RunE(cmd, nil); err != nil { + t.Fatalf("AgentCmd.RunE() error = %v", err) + } + if got := stdout.String(); got != "answer\n" { + t.Errorf("stdout = %q, want print answer", got) + } + if strings.Contains(stdout.String(), "\x1b") || strings.Contains(stderr.String(), "\x1b") { + t.Errorf("print output contains escape bytes: stdout=%q stderr=%q", stdout.String(), stderr.String()) + } +} + +func TestRunAgentTUIPrintsTranscriptAfterRun(t *testing.T) { + model := aiagent.NewTUIModel(context.Background(), nil, aiagent.TUIOptions{NoColor: true}) + runErr := errors.New("run failed") + var output bytes.Buffer + err := runAgentTUI(&output, &output, func() (tea.Model, error) { + _, _ = output.WriteString("after run\n") + return model, runErr + }) + if !errors.Is(err, runErr) { + t.Errorf("runAgentTUI() error = %v, want %v", err, runErr) + } + got := output.String() + if !strings.HasPrefix(got, "after run\n") { + t.Fatalf("output = %q, want run output first", got) + } + if !strings.Contains(got, "New session") { + t.Errorf("redirected output = %q, want final transcript after Run", got) + } + if count := strings.Count(got, "New session"); count != 1 { + t.Errorf("transcript header count = %d, want 1", count) + } +} + +func TestRunAgentTUIWritesPlainStdoutAndPreservesUITranscript(t *testing.T) { + model := aiagent.NewTUIModel(context.Background(), nil, aiagent.TUIOptions{Version: "1.2.3"}) + wantTranscript := model.Transcript() + if !strings.Contains(wantTranscript, "\x1b") { + t.Fatal("test transcript has no ANSI styling") + } + var stdout, uiWriter bytes.Buffer + if err := runAgentTUI(&stdout, &uiWriter, func() (tea.Model, error) { + return model, nil + }); err != nil { + t.Fatalf("runAgentTUI() error = %v", err) + } + if strings.Contains(stdout.String(), "\x1b") { + t.Errorf("stdout transcript = %q, want no ANSI escape bytes", stdout.String()) + } + if got, want := stdout.String(), ansi.Strip(wantTranscript); got != want { + t.Errorf("stdout transcript = %q, want stripped transcript %q", got, want) + } + if got := uiWriter.String(); got != wantTranscript { + t.Errorf("UI transcript = %q, want unchanged transcript %q", got, wantTranscript) + } + if !strings.Contains(uiWriter.String(), "New session") { + t.Errorf("UI writer = %q, want final transcript", uiWriter.String()) + } +} + +type failingTranscriptWriter struct { + err error +} + +func (w failingTranscriptWriter) Write([]byte) (int, error) { + return 0, w.err +} + +func TestRunAgentTUITranscriptWriteErrorPrecedence(t *testing.T) { + model := aiagent.NewTUIModel(context.Background(), nil, aiagent.TUIOptions{NoColor: true}) + runErr := errors.New("run failed") + stdoutWriteErr := errors.New("stdout write failed") + uiWriteErr := errors.New("UI write failed") + + t.Run("run error takes precedence", func(t *testing.T) { + err := runAgentTUI( + failingTranscriptWriter{err: stdoutWriteErr}, + failingTranscriptWriter{err: uiWriteErr}, + func() (tea.Model, error) { return model, runErr }, + ) + if !errors.Is(err, runErr) { + t.Errorf("runAgentTUI() error = %v, want %v", err, runErr) + } + }) + + t.Run("first write error takes precedence", func(t *testing.T) { + err := runAgentTUI( + failingTranscriptWriter{err: stdoutWriteErr}, + failingTranscriptWriter{err: uiWriteErr}, + func() (tea.Model, error) { return model, nil }, + ) + if !errors.Is(err, stdoutWriteErr) { + t.Errorf("runAgentTUI() error = %v, want %v", err, stdoutWriteErr) + } + }) + + t.Run("UI write error returned when stdout succeeds", func(t *testing.T) { + var stdout bytes.Buffer + err := runAgentTUI( + &stdout, + failingTranscriptWriter{err: uiWriteErr}, + func() (tea.Model, error) { return model, nil }, + ) + if !errors.Is(err, uiWriteErr) { + t.Errorf("runAgentTUI() error = %v, want %v", err, uiWriteErr) + } + }) +} + +func TestPrintModeExitsTwoWhenNotAuthenticated(t *testing.T) { + t.Setenv("HOME", t.TempDir()) + previousKey := viper.GetString("api_key") + viper.Set("api_key", "") + t.Cleanup(func() { viper.Set("api_key", previousKey) }) + cmd := &cobra.Command{} + cmd.Flags().String("print", "", "") + if err := cmd.Flags().Set("print", "hello"); err != nil { + t.Fatal(err) + } + cmd.SetContext(context.Background()) + _, err := getAPIClient() + if !aiagent.IsNotAuthenticated(err) { + t.Fatalf("getAPIClient() error = %v, want not authenticated", err) + } + var coder interface{ ExitCode() int } + if !isPrintMode(cmd) { + t.Fatal("isPrintMode() = false with --print") + } + if wrapped := aiagent.PrintExitError(err); !errors.As(wrapped, &coder) || coder.ExitCode() != 2 { + t.Errorf("print exit error = %v, want exit code 2", wrapped) + } + if !strings.Contains(AgentCmd.Long, "2 not authenticated") { + t.Errorf("agent help = %q, want exit codes documented", AgentCmd.Long) + } +} diff --git a/internal/cmd/agent/flags_test.go b/internal/cmd/agent/flags_test.go new file mode 100644 index 0000000..fbfb410 --- /dev/null +++ b/internal/cmd/agent/flags_test.go @@ -0,0 +1,227 @@ +package agent + +import ( + "bytes" + "context" + "net/http" + "net/http/httptest" + "os" + "path/filepath" + "strings" + "sync/atomic" + "testing" + "time" + + "github.com/spf13/cobra" + "github.com/spf13/viper" + + aiagent "github.com/rootlyhq/rootly-cli/internal/agent" +) + +func agentFlagsCommand(t *testing.T, flags map[string]string) *cobra.Command { + t.Helper() + cmd := &cobra.Command{} + cmd.Flags().String("print", "", "") + cmd.Flags().String("resume", "", "") + cmd.Flags().Bool("continue", false, "") + cmd.Flags().String("incident", "", "") + cmd.Flags().String("alert", "", "") + cmd.Flags().String("output-format", "text", "") + for name, value := range flags { + if err := cmd.Flags().Set(name, value); err != nil { + t.Fatal(err) + } + } + cmd.SetIn(strings.NewReader("")) + cmd.SetContext(context.Background()) + return cmd +} + +func TestRunAgentRejectsBadFlagsBeforeAnyRequest(t *testing.T) { + t.Setenv("HOME", t.TempDir()) + previous := viper.GetString("api_key") + viper.Set("api_key", "test-key") + viper.Set("api_host", "127.0.0.1:1") + t.Cleanup(func() { viper.Set("api_key", previous); viper.Set("api_host", "") }) + cases := []struct { + flags map[string]string + want string + }{ + {map[string]string{"print": "hi", "output-format": "yaml"}, `invalid output format "yaml"`}, + {map[string]string{"print": "hi", "resume": "not-a-uuid"}, `invalid resume session UUID "not-a-uuid"`}, + } + for _, tc := range cases { + err := runAgent(agentFlagsCommand(t, tc.flags)) + if err == nil || !strings.Contains(err.Error(), tc.want) { + t.Errorf("runAgent(%v) error = %v, want %q", tc.flags, err, tc.want) + } + } +} + +func TestRunAgentNeedsAuthenticationInPrintMode(t *testing.T) { + t.Setenv("HOME", t.TempDir()) + previous := viper.GetString("api_key") + viper.Set("api_key", "") + t.Cleanup(func() { viper.Set("api_key", previous) }) + err := runAgent(agentFlagsCommand(t, map[string]string{"print": "hi"})) + if err == nil || !strings.Contains(err.Error(), "authentication required") { + t.Fatalf("runAgent without credentials = %v", err) + } +} + +func TestAgentWorkingDir(t *testing.T) { + dir := t.TempDir() + t.Chdir(dir) + got := agentWorkingDir() + resolvedGot, _ := os.Stat(got) + resolvedWant, _ := os.Stat(dir) + if !os.SameFile(resolvedGot, resolvedWant) { + t.Fatalf("agentWorkingDir() = %q, want %q", got, dir) + } +} + +func TestRunAgentReportsLookupFailures(t *testing.T) { + t.Setenv("HOME", t.TempDir()) + unauthorized := false + server := httptest.NewServer(http.HandlerFunc(func(writer http.ResponseWriter, request *http.Request) { + if unauthorized { + writer.WriteHeader(http.StatusUnauthorized) + return + } + http.NotFound(writer, request) + })) + defer server.Close() + oldKey, oldHost := viper.GetString("api_key"), viper.GetString("api_host") + viper.Set("api_key", "test-key") + viper.Set("api_host", server.URL) + t.Cleanup(func() { viper.Set("api_key", oldKey); viper.Set("api_host", oldHost) }) + cases := []struct { + flags map[string]string + want string + }{ + {map[string]string{"print": "hi", "incident": "INC-404"}, "incident not found"}, + {map[string]string{"print": "hi", "alert": "missing"}, "alert not found"}, + {map[string]string{"print": "hi", "continue": "true"}, "404"}, + } + for _, tc := range cases { + err := runAgent(agentFlagsCommand(t, tc.flags)) + if err == nil || !strings.Contains(err.Error(), tc.want) { + t.Errorf("runAgent(%v) error = %v, want %q", tc.flags, err, tc.want) + } + } + unauthorized = true + err := runAgent(agentFlagsCommand(t, map[string]string{"print": "hi", "incident": "INC-1"})) + if !aiagent.IsNotAuthenticated(err) { + t.Fatalf("a 401 preflight must stop as not authenticated, got %v", err) + } +} + +func TestRunAgentPrintModeStartsLocalToolsAndSurfacesStreamErrors(t *testing.T) { + t.Setenv("HOME", t.TempDir()) + t.Chdir(t.TempDir()) + var streamed atomic.Int32 + server := httptest.NewServer(http.HandlerFunc(func(writer http.ResponseWriter, request *http.Request) { + if strings.Contains(request.URL.Path, "/ai/") { + streamed.Add(1) + writer.WriteHeader(http.StatusInternalServerError) + return + } + http.NotFound(writer, request) + })) + defer server.Close() + oldKey, oldHost := viper.GetString("api_key"), viper.GetString("api_host") + viper.Set("api_key", "test-key") + viper.Set("api_host", server.URL) + t.Cleanup(func() { viper.Set("api_key", oldKey); viper.Set("api_host", oldHost) }) + cmd := agentFlagsCommand(t, map[string]string{"print": "hello"}) + cmd.AddCommand(&cobra.Command{Use: "child"}) + var stdout, stderr bytes.Buffer + cmd.SetOut(&stdout) + cmd.SetErr(&stderr) + err := runAgent(cmd) + if err == nil { + t.Fatal("a failing chat stream must return an error") + } + if streamed.Load() == 0 { + t.Fatalf("print mode never reached the chat API: %v", err) + } + if strings.Contains(stdout.String(), "\x1b[2J") { + t.Fatal("print mode must not clear the terminal") + } +} + +func TestPreflightAIDebugReportsSoftFailures(t *testing.T) { + var status atomic.Int32 + server := httptest.NewServer(http.HandlerFunc(func(writer http.ResponseWriter, _ *http.Request) { + writer.Header().Set("Content-Type", "application/vnd.api+json") + writer.WriteHeader(int(status.Load())) + _, _ = writer.Write([]byte(`{"errors":[{"title":"Forbidden"}]}`)) + })) + oldKey, oldHost, oldDebug := viper.GetString("api_key"), viper.GetString("api_host"), viper.GetBool("debug") + viper.Set("api_key", "test-key") + viper.Set("api_host", server.URL) + viper.Set("debug", true) + t.Cleanup(func() { viper.Set("api_key", oldKey); viper.Set("api_host", oldHost); viper.Set("debug", oldDebug) }) + client, err := getAPIClient() + if err != nil { + t.Fatal(err) + } + for _, code := range []int32{http.StatusForbidden, http.StatusBadGateway} { + status.Store(code) + cmd := agentFlagsCommand(t, nil) + var stderr bytes.Buffer + cmd.SetErr(&stderr) + user, err := preflightAI(cmd, client) + if user != nil || err != nil { + t.Fatalf("HTTP %d preflight = %v, %v; want a soft failure", code, user, err) + } + if !strings.Contains(stderr.String(), "AI preflight returned HTTP") { + t.Errorf("HTTP %d debug output = %q", code, stderr.String()) + } + } + server.Close() + cmd := agentFlagsCommand(t, nil) + var stderr bytes.Buffer + cmd.SetErr(&stderr) + if user, err := preflightAI(cmd, client); user != nil || err != nil || !strings.Contains(stderr.String(), "AI preflight failed") { + t.Fatalf("unreachable preflight = %v, %v, %q", user, err, stderr.String()) + } +} + +func TestGetAPIClientDefaultsTheEndpoint(t *testing.T) { + oldKey, oldHost := viper.GetString("api_key"), viper.GetString("api_host") + viper.Set("api_key", "test-key") + viper.Set("api_host", "") + t.Cleanup(func() { viper.Set("api_key", oldKey); viper.Set("api_host", oldHost) }) + client, err := getAPIClient() + if err != nil || client.APIHost() == "" { + t.Fatalf("getAPIClient() = %v, %v", client, err) + } +} + +func TestStartLocalMCPReportsBadConfigAndKeepsGoing(t *testing.T) { + home := t.TempDir() + t.Setenv("HOME", home) + t.Setenv("USERPROFILE", home) + t.Chdir(t.TempDir()) + if err := os.MkdirAll(filepath.Join(home, ".rootly-cli"), 0o700); err != nil { + t.Fatal(err) + } + if err := os.WriteFile(filepath.Join(home, ".rootly-cli", "mcp.json"), []byte("{not json"), 0o600); err != nil { + t.Fatal(err) + } + cmd := agentFlagsCommand(t, nil) + var stderr bytes.Buffer + cmd.SetErr(&stderr) + manager := startLocalMCP(cmd) + defer manager.Close() + if !strings.Contains(stderr.String(), "local MCP config") { + t.Fatalf("bad config not reported: %q", stderr.String()) + } + ctx, cancel := context.WithCancel(context.Background()) + cancel() + waitForLocalMCP(ctx, manager, time.Second) + if tools, _ := manager.ClientTools(); tools == nil && manager.Connecting() != 0 { + t.Fatal("manager should still be usable after a bad config") + } +} diff --git a/internal/cmd/agent_register.go b/internal/cmd/agent_register.go new file mode 100644 index 0000000..6a56a49 --- /dev/null +++ b/internal/cmd/agent_register.go @@ -0,0 +1,7 @@ +package cmd + +import "github.com/rootlyhq/rootly-cli/internal/cmd/agent" + +func init() { + rootCmd.AddCommand(agent.AgentCmd) +} diff --git a/internal/cmd/agent_register_test.go b/internal/cmd/agent_register_test.go new file mode 100644 index 0000000..c286f7e --- /dev/null +++ b/internal/cmd/agent_register_test.go @@ -0,0 +1,36 @@ +package cmd + +import ( + "bytes" + "strings" + "testing" +) + +func TestAgentAliasHelpAndHiddenRootHelp(t *testing.T) { + var stdout, stderr bytes.Buffer + rootCmd.SetOut(&stdout) + rootCmd.SetErr(&stderr) + defer func() { + rootCmd.SetOut(nil) + rootCmd.SetErr(nil) + rootCmd.SetArgs([]string{}) + }() + + rootCmd.SetArgs([]string{"chat", "--help"}) + if err := rootCmd.Execute(); err != nil { + t.Fatalf("rootly chat --help returned error: %v", err) + } + if !strings.Contains(stdout.String(), "rootly agent") { + t.Errorf("chat help does not show agent command: %q", stdout.String()) + } + + stdout.Reset() + stderr.Reset() + rootCmd.SetArgs([]string{"--help"}) + if err := rootCmd.Execute(); err != nil { + t.Fatalf("rootly --help returned error: %v", err) + } + if strings.Contains(stdout.String(), "rootly agent") || strings.Contains(stdout.String(), "rootly chat") { + t.Errorf("root help lists hidden agent command: %q", stdout.String()) + } +} diff --git a/internal/cmd/auth/cmd_test.go b/internal/cmd/auth/cmd_test.go index 3f748e3..64d6d4a 100644 --- a/internal/cmd/auth/cmd_test.go +++ b/internal/cmd/auth/cmd_test.go @@ -6,6 +6,7 @@ import ( "testing" "time" + "github.com/rootlyhq/rootly-cli/internal/config" "github.com/rootlyhq/rootly-cli/internal/oauth" ) @@ -13,6 +14,15 @@ func TestLogoutCmd(t *testing.T) { tmpDir := t.TempDir() t.Setenv("HOME", tmpDir) t.Setenv("USERPROFILE", tmpDir) + t.Setenv("ROOTLY_API_KEY", "") + + if err := config.Save(&config.Config{ + APIKey: "test-api-key", + ClientID: "stale-client", + Scopes: []string{"openid", "ai.chat:read"}, + }); err != nil { + t.Fatalf("failed to write config: %v", err) + } // Save tokens first tokens := &oauth.TokenData{ @@ -36,6 +46,19 @@ func TestLogoutCmd(t *testing.T) { if _, err := oauth.LoadTokens(); !os.IsNotExist(err) { t.Errorf("expected tokens to be cleared, got err: %v", err) } + cfg, err := config.Load() + if err != nil { + t.Fatalf("failed to load config after logout: %v", err) + } + if cfg.ClientID != "" { + t.Errorf("client_id = %q, want empty", cfg.ClientID) + } + if len(cfg.Scopes) != 0 { + t.Errorf("scopes = %v, want empty", cfg.Scopes) + } + if cfg.APIKey != "test-api-key" { + t.Errorf("api_key = %q, want preserved API key", cfg.APIKey) + } } func TestLogoutCmd_NoTokens(t *testing.T) { diff --git a/internal/cmd/auth/logout.go b/internal/cmd/auth/logout.go index 368d110..31413f5 100644 --- a/internal/cmd/auth/logout.go +++ b/internal/cmd/auth/logout.go @@ -10,13 +10,16 @@ import ( var LogoutCmd = &cobra.Command{ Use: "logout", - Short: "Clear stored OAuth2 tokens", + Short: "Clear stored OAuth2 tokens and client registration", Annotations: map[string]string{"skipAuth": "true"}, RunE: func(cmd *cobra.Command, args []string) error { if err := oauth.ClearTokens(); err != nil { return fmt.Errorf("failed to clear tokens: %w", err) } - _, _ = fmt.Fprintf(cmd.OutOrStderr(), "Logged out. OAuth tokens cleared.\n") + if err := oauth.ClearRegistration(); err != nil { + return fmt.Errorf("failed to clear OAuth client registration: %w", err) + } + _, _ = fmt.Fprintf(cmd.OutOrStderr(), "Logged out. OAuth tokens and client registration cleared.\n") return nil }, } diff --git a/internal/cmd/exit_code_test.go b/internal/cmd/exit_code_test.go new file mode 100644 index 0000000..5830f48 --- /dev/null +++ b/internal/cmd/exit_code_test.go @@ -0,0 +1,28 @@ +package cmd + +import ( + "errors" + "fmt" + "testing" + + aiagent "github.com/rootlyhq/rootly-cli/internal/agent" +) + +func TestExitCode(t *testing.T) { + notAuthenticated := aiagent.PrintExitError(aiagent.NotAuthenticatedError("not authenticated")) + tests := []struct { + name string + err error + want int + }{ + {name: "generic", err: errors.New("boom"), want: 1}, + {name: "print not authenticated", err: notAuthenticated, want: 2}, + {name: "wrapped", err: fmt.Errorf("agent: %w", notAuthenticated), want: 2}, + {name: "unauthenticated outside print mode", err: aiagent.NotAuthenticatedError("not authenticated"), want: 1}, + } + for _, test := range tests { + if got := exitCode(test.err); got != test.want { + t.Errorf("%s: exitCode() = %d, want %d", test.name, got, test.want) + } + } +} diff --git a/internal/cmd/helpers_test.go b/internal/cmd/helpers_test.go index b40d642..260fbfc 100644 --- a/internal/cmd/helpers_test.go +++ b/internal/cmd/helpers_test.go @@ -173,8 +173,10 @@ func TestErrAborted(t *testing.T) { func TestSetVersionInfo(t *testing.T) { // Save originals origV, origC, origD := version, commit, date + origRootVersion := rootCmd.Version defer func() { version, commit, date = origV, origC, origD + rootCmd.Version = origRootVersion }() SetVersionInfo("1.2.3", "abc123", "2025-06-15") @@ -188,14 +190,19 @@ func TestSetVersionInfo(t *testing.T) { if date != "2025-06-15" { t.Errorf("date = %q, want %q", date, "2025-06-15") } + if rootCmd.Version != "1.2.3" { + t.Errorf("root command version = %q, want %q", rootCmd.Version, "1.2.3") + } } // --- version command tests --- func TestVersionCommandText(t *testing.T) { origV, origC, origD := version, commit, date + origRootVersion := rootCmd.Version defer func() { version, commit, date = origV, origC, origD + rootCmd.Version = origRootVersion }() SetVersionInfo("2.0.0", "def456", "2025-07-01") @@ -224,8 +231,10 @@ func TestVersionCommandText(t *testing.T) { func TestVersionCommandJSON(t *testing.T) { origV, origC, origD := version, commit, date + origRootVersion := rootCmd.Version defer func() { version, commit, date = origV, origC, origD + rootCmd.Version = origRootVersion }() SetVersionInfo("3.0.0", "abc789", "2025-08-01") diff --git a/internal/cmd/mcp/mcp.go b/internal/cmd/mcp/mcp.go new file mode 100644 index 0000000..34e1192 --- /dev/null +++ b/internal/cmd/mcp/mcp.go @@ -0,0 +1,623 @@ +// Package mcp implements `rootly mcp`, which manages local MCP servers for +// `rootly agent` using the same file shapes as Claude Code. +package mcp + +import ( + "bufio" + "context" + "encoding/json" + "errors" + "fmt" + "io" + "net/http" + "os" + "path/filepath" + "sort" + "strconv" + "strings" + "time" + + "github.com/mattn/go-isatty" + "github.com/spf13/cobra" + + localmcp "github.com/rootlyhq/rootly-cli/internal/mcp" +) + +var MCPCmd = &cobra.Command{ + Use: "mcp", + Short: "Manage local MCP servers for rootly agent", + Hidden: true, + Annotations: map[string]string{"skipAuth": "true"}, + Args: cobra.NoArgs, + RunE: func(cmd *cobra.Command, _ []string) error { + return cmd.Help() + }, +} + +// Environment is swappable in tests. +type Environment struct { + CWD func() (string, error) + Store func(cwd string) *localmcp.Store + State func() *localmcp.StateStore + Tokens func() *localmcp.TokenStore + Discovery func(cwd string) localmcp.Discovery + IsTTY func() bool + Stdin io.Reader + HTTP *http.Client + OpenURL func(string) error +} + +var env = DefaultEnvironment() + +func DefaultEnvironment() Environment { + return Environment{ + CWD: os.Getwd, + Store: localmcp.NewStore, + State: localmcp.NewStateStore, + Tokens: localmcp.NewTokenStore, + Discovery: localmcp.NewDiscovery, + IsTTY: func() bool { return isatty.IsTerminal(os.Stdin.Fd()) }, + Stdin: os.Stdin, + HTTP: &http.Client{Timeout: 5 * time.Second}, + OpenURL: localmcp.OpenBrowser, + } +} + +// SetEnvironment replaces the environment and returns a restore func. +func SetEnvironment(next Environment) func() { + previous := env + env = next + return func() { env = previous } +} + +func cwd() (string, error) { + dir, err := env.CWD() + if err != nil { + return "", err + } + return filepath.Abs(dir) +} + +func displayPath(path string) string { + home, err := os.UserHomeDir() + if err == nil && home != "" && strings.HasPrefix(path, home+string(os.PathSeparator)) { + return "~" + strings.TrimPrefix(path, home) + } + return path +} + +func scopeTarget(store *localmcp.Store, scope string) string { + switch scope { + case localmcp.ScopeProject: + return "project config (" + displayPath(store.ProjectPath()) + ")" + case localmcp.ScopeLocal: + return "local config (" + displayPath(store.UserPath) + " [project: " + store.CWD + "])" + default: + return "user config (" + displayPath(store.UserPath) + ")" + } +} + +func validScope(scope string) error { + switch scope { + case localmcp.ScopeLocal, localmcp.ScopeUser, localmcp.ScopeProject: + return nil + } + return fmt.Errorf("invalid scope %q: expected local, user or project", scope) +} + +var ( + addTransport string + addScope string + addEnv []string + addHeaders []string + addURL string + addBearerEnv string + addJSONScope string + removeScope string + importAll bool +) + +var addCmd = &cobra.Command{ + Use: "add [flags] [args...]", + Short: "Add an MCP server (same syntax as `claude mcp add`)", + Example: ` rootly mcp add --transport http braintrust https://api.braintrust.dev/mcp + rootly mcp add --transport http sentry https://mcp.sentry.dev/mcp -H "Authorization: Bearer ${SENTRY_TOKEN}" + rootly mcp add scalr -e SCALR_API_URL=https://rootly.scalr.io -- docker run --rm -i -e SCALR_API_TOKEN scalr/mcp-server + rootly mcp add -s project maestro -- maestro mcp + rootly mcp add linear --url https://mcp.linear.app/mcp --bearer-token-env-var LINEAR_TOKEN`, + Args: cobra.MinimumNArgs(1), + RunE: func(cmd *cobra.Command, args []string) error { + server, err := buildAddConfig(args) + if err != nil { + return err + } + return addServer(cmd, args[0], addScope, server) + }, +} + +func buildAddConfig(args []string) (localmcp.ServerConfig, error) { + transport := strings.ToLower(addTransport) + positional := args[1:] + url := addURL + if url == "" && transport == "" && len(positional) == 1 && looksLikeURL(positional[0]) { + transport = localmcp.TransportHTTP + } + if url != "" { + if len(positional) > 0 { + return localmcp.ServerConfig{}, errors.New("--url takes the URL; don't also pass a command or URL argument") + } + if transport == "" { + transport = localmcp.TransportHTTP + } + } + if transport == "" { + transport = localmcp.TransportStdio + } + envs, err := parsePairs(addEnv, "=", "-e/--env", "KEY=value") + if err != nil { + return localmcp.ServerConfig{}, err + } + headers, err := parsePairs(addHeaders, ":", "-H/--header", `"Name: value"`) + if err != nil { + return localmcp.ServerConfig{}, err + } + switch transport { + case localmcp.TransportStdio: + if len(positional) == 0 { + return localmcp.ServerConfig{}, errors.New("missing command: rootly mcp add -- [args...]") + } + if len(headers) > 0 || addBearerEnv != "" { + return localmcp.ServerConfig{}, errors.New("-H/--header and --bearer-token-env-var only apply to http and sse servers") + } + return localmcp.ServerConfig{Command: positional[0], Args: positional[1:], Env: envs}, nil + case localmcp.TransportHTTP, localmcp.TransportSSE: + if url == "" { + if len(positional) != 1 { + return localmcp.ServerConfig{}, fmt.Errorf("%s servers take exactly one URL: rootly mcp add --transport %s ", transport, transport) + } + url = positional[0] + } + if !looksLikeURL(url) { + return localmcp.ServerConfig{}, fmt.Errorf("invalid URL %q: expected http:// or https://", url) + } + if len(envs) > 0 { + return localmcp.ServerConfig{}, errors.New("-e/--env only applies to stdio servers") + } + return localmcp.ServerConfig{Type: transport, URL: url, Headers: headers, BearerTokenEnvVar: addBearerEnv}, nil + default: + return localmcp.ServerConfig{}, fmt.Errorf("invalid transport %q: expected stdio, http or sse", addTransport) + } +} + +func looksLikeURL(value string) bool { + lower := strings.ToLower(value) + return strings.HasPrefix(lower, "http://") || strings.HasPrefix(lower, "https://") +} + +func parsePairs(values []string, separator, flag, shape string) (map[string]string, error) { + if len(values) == 0 { + return nil, nil + } + pairs := make(map[string]string, len(values)) + for _, value := range values { + key, rest, ok := strings.Cut(value, separator) + key = strings.TrimSpace(key) + if !ok || key == "" { + return nil, fmt.Errorf("invalid %s value %q: expected %s", flag, maskedValue(value, separator), shape) + } + if separator == ":" { + rest = strings.TrimSpace(rest) + } + pairs[key] = rest + } + return pairs, nil +} + +func maskedValue(value, separator string) string { + if key, _, ok := strings.Cut(value, separator); ok { + return key + separator + localmcp.Mask + } + return value +} + +func addServer(cmd *cobra.Command, name, scope string, server localmcp.ServerConfig) error { + if err := validScope(scope); err != nil { + return err + } + dir, err := cwd() + if err != nil { + return err + } + store := env.Store(dir) + if err := store.Add(scope, name, server); err != nil { + return err + } + out := cmd.OutOrStdout() + kind := strings.ToUpper(server.Transport()) + if server.Transport() == localmcp.TransportStdio { + kind = "stdio" + } + _, _ = fmt.Fprintf(out, "Added %s MCP server %s to %s\n", kind, name, scopeTarget(store, scope)) + if scope == localmcp.ScopeProject { + _, _ = fmt.Fprintln(out, " project servers need approval before they start: open /mcp in rootly agent") + } + if server.Transport() != localmcp.TransportStdio && needsAuthentication(cmd.Context(), server) { + _, _ = fmt.Fprintf(out, " needs authentication: run `rootly mcp login %s` or /mcp\n", name) + } + return nil +} + +func needsAuthentication(ctx context.Context, server localmcp.ServerConfig) bool { + if ctx == nil { + ctx = context.Background() + } + if server.BearerTokenEnvVar != "" { + return false + } + for key := range server.Headers { + if strings.EqualFold(key, "Authorization") { + return false + } + } + if _, ok := env.Tokens().Get(server.URL); ok { + return false + } + ctx, cancel := context.WithTimeout(ctx, 5*time.Second) + defer cancel() + status, err := localmcp.ProbeStatus(ctx, env.HTTP, server.Expanded(nil)) + return err == nil && status == http.StatusUnauthorized +} + +var addJSONCmd = &cobra.Command{ + Use: "add-json ", + Short: "Add an MCP server from a JSON config", + Example: ` rootly mcp add-json dart '{"command":"fvm","args":["dart","mcp-server"]}'`, + Args: cobra.ExactArgs(2), + RunE: func(cmd *cobra.Command, args []string) error { + var server localmcp.ServerConfig + decoder := json.NewDecoder(strings.NewReader(args[1])) + decoder.DisallowUnknownFields() + if err := decoder.Decode(&server); err != nil { + return fmt.Errorf("invalid MCP server JSON: %w", err) + } + return addServer(cmd, args[0], addJSONScope, server) + }, +} + +type resolved struct { + server localmcp.Server + store *localmcp.Store +} + +func allServers() ([]localmcp.Server, *localmcp.Store, error) { + dir, err := cwd() + if err != nil { + return nil, nil, err + } + store := env.Store(dir) + servers, err := store.All() + imported, _ := env.Discovery(dir).Imported() + return append(servers, imported...), store, err +} + +func findServer(name, scope string) (resolved, error) { + servers, store, err := allServers() + if err != nil && len(servers) == 0 { + return resolved{}, err + } + var matches []localmcp.Server + for _, server := range servers { + if server.Name == name && (scope == "" || server.Scope == scope) { + matches = append(matches, server) + } + } + if len(matches) == 0 { + if scope != "" { + return resolved{}, fmt.Errorf("no MCP server named %q in %s config", name, scope) + } + return resolved{}, fmt.Errorf("no MCP server named %q (see `rootly mcp list`)", name) + } + return resolved{server: matches[0], store: store}, nil +} + +var listCmd = &cobra.Command{ + Use: "list", + Aliases: []string{"ls"}, + Short: "List configured and imported MCP servers", + Args: cobra.NoArgs, + RunE: func(cmd *cobra.Command, _ []string) error { + servers, _, err := allServers() + if err != nil && len(servers) == 0 { + return err + } + out := cmd.OutOrStdout() + if len(servers) == 0 { + _, _ = fmt.Fprintln(out, "No MCP servers configured. Add one with `rootly mcp add`, or copy existing ones with `rootly mcp import`.") + return nil + } + state, _ := env.State().Load() + tokens := env.Tokens() + for _, server := range servers { + _, _ = fmt.Fprintf(out, "%s: %s (%s) - %s\n", server.Name, server.Config.Origin(), server.Config.Transport(), describeState(state, tokens, server)) + } + return err + }, +} + +func describeState(state localmcp.State, tokens *localmcp.TokenStore, server localmcp.Server) string { + scope := server.Scope + if server.Scope == localmcp.ScopeImported { + scope = "imported from " + server.SourceLabel + ", read-only" + } + parts := []string{scope} + switch { + case !state.IsEnabled(server): + parts = append(parts, "off") + case state.NeedsApproval(server): + parts = append(parts, "pending approval") + } + if server.Config.Transport() != localmcp.TransportStdio { + if _, ok := tokens.Get(server.Config.URL); ok { + parts = append(parts, "logged in") + } + } + return strings.Join(parts, ", ") +} + +var getCmd = &cobra.Command{ + Use: "get ", + Short: "Show an MCP server's config (secrets masked)", + Args: cobra.ExactArgs(1), + RunE: func(cmd *cobra.Command, args []string) error { + found, err := findServer(args[0], "") + if err != nil { + return err + } + server := found.server + config := server.Config.Masked() + state, _ := env.State().Load() + out := cmd.OutOrStdout() + _, _ = fmt.Fprintf(out, "%s:\n", server.Name) + _, _ = fmt.Fprintf(out, " Scope: %s\n", describeState(state, env.Tokens(), server)) + _, _ = fmt.Fprintf(out, " Source: %s\n", displayPath(server.Path)) + _, _ = fmt.Fprintf(out, " Type: %s\n", config.Transport()) + if config.Transport() == localmcp.TransportStdio { + _, _ = fmt.Fprintf(out, " Command: %s\n", config.Command) + _, _ = fmt.Fprintf(out, " Args: %s\n", strings.Join(config.Args, " ")) + if len(config.Env) > 0 { + _, _ = fmt.Fprintf(out, " Environment: %s\n", strings.Join(localmcp.MaskedPairs(config.Env, "="), ", ")) + } + } else { + _, _ = fmt.Fprintf(out, " URL: %s\n", config.URL) + if len(config.Headers) > 0 { + _, _ = fmt.Fprintf(out, " Headers: %s\n", strings.Join(localmcp.MaskedPairs(config.Headers, ": "), ", ")) + } + if config.BearerTokenEnvVar != "" { + _, _ = fmt.Fprintf(out, " Bearer token env var: %s\n", config.BearerTokenEnvVar) + } + } + if server.Scope != localmcp.ScopeImported { + _, _ = fmt.Fprintf(out, "\nTo remove this server, run: rootly mcp remove %s -s %s\n", server.Name, server.Scope) + } + return nil + }, +} + +var removeCmd = &cobra.Command{ + Use: "remove ", + Aliases: []string{"rm"}, + Short: "Remove an MCP server", + Args: cobra.ExactArgs(1), + RunE: func(cmd *cobra.Command, args []string) error { + dir, err := cwd() + if err != nil { + return err + } + store := env.Store(dir) + scopes := []string{removeScope} + if removeScope == "" { + scopes = nil + for _, scope := range []string{localmcp.ScopeLocal, localmcp.ScopeProject, localmcp.ScopeUser} { + servers, _ := store.Servers(scope) + for _, server := range servers { + if server.Name == args[0] { + scopes = append(scopes, scope) + } + } + } + switch len(scopes) { + case 0: + if found, findErr := findServer(args[0], localmcp.ScopeImported); findErr == nil { + return fmt.Errorf("%s is imported from %s (read-only); edit %s instead", args[0], found.server.SourceLabel, displayPath(found.server.Path)) + } + return fmt.Errorf("no MCP server named %q", args[0]) + case 1: + default: + return fmt.Errorf("MCP server %q exists in several scopes (%s); pass -s/--scope", args[0], strings.Join(scopes, ", ")) + } + } else if err := validScope(removeScope); err != nil { + return err + } + if err := store.Remove(scopes[0], args[0]); err != nil { + return err + } + _, _ = fmt.Fprintf(cmd.OutOrStdout(), "Removed MCP server %s from %s\n", args[0], scopeTarget(store, scopes[0])) + return nil + }, +} + +var loginCmd = &cobra.Command{ + Use: "login ", + Short: "Log in to an HTTP MCP server with OAuth", + Args: cobra.ExactArgs(1), + RunE: func(cmd *cobra.Command, args []string) error { + found, err := findServer(args[0], "") + if err != nil { + return err + } + config := found.server.Config.Expanded(nil) + if config.Transport() == localmcp.TransportStdio { + return fmt.Errorf("%s is a stdio server; OAuth login only applies to http and sse servers", args[0]) + } + errOut := cmd.ErrOrStderr() + ctx, cancel := context.WithTimeout(cmd.Context(), 5*time.Minute) + defer cancel() + err = env.Tokens().Login(ctx, config, localmcp.LoginOptions{OpenURL: func(url string) { + _, _ = fmt.Fprintf(errOut, "Opening your browser to authorize %s.\nIf it doesn't open, visit:\n %s\n", args[0], url) + _ = env.OpenURL(url) + }}) + if err != nil { + return fmt.Errorf("login to %s failed: %w", args[0], err) + } + _, _ = fmt.Fprintf(cmd.OutOrStdout(), "Logged in to %s. Tokens saved in ~/.rootly-cli (0600).\n", args[0]) + return nil + }, +} + +var logoutCmd = &cobra.Command{ + Use: "logout ", + Short: "Forget the OAuth tokens of an HTTP MCP server", + Args: cobra.ExactArgs(1), + RunE: func(cmd *cobra.Command, args []string) error { + found, err := findServer(args[0], "") + if err != nil { + return err + } + removed, err := env.Tokens().Delete(found.server.Config.Expanded(nil).URL) + if err != nil { + return err + } + if !removed { + _, _ = fmt.Fprintf(cmd.OutOrStdout(), "Not logged in to %s.\n", args[0]) + return nil + } + _, _ = fmt.Fprintf(cmd.OutOrStdout(), "Logged out of %s.\n", args[0]) + return nil + }, +} + +var importCmd = &cobra.Command{ + Use: "import [names...]", + Short: "Copy MCP servers from Claude Code, Codex, Cursor, VS Code or Claude Desktop into your user config", + Long: `Copy MCP servers from Claude Code (~/.claude.json), Codex (~/.codex/config.toml), +Cursor (~/.cursor/mcp.json, ./.cursor/mcp.json), VS Code (./.vscode/mcp.json) and +Claude Desktop into ~/.rootly-cli/mcp.json. Without names it asks which ones to copy.`, + RunE: func(cmd *cobra.Command, args []string) error { + dir, err := cwd() + if err != nil { + return err + } + store := env.Store(dir) + discovered, problems := env.Discovery(dir).Imported() + errOut := cmd.ErrOrStderr() + for _, problem := range problems { + _, _ = fmt.Fprintf(errOut, "warning: could not read %s (%s): %v\n", problem.Source.Label, displayPath(problem.Source.Path), problem.Err) + } + if len(discovered) == 0 { + _, _ = fmt.Fprintln(cmd.OutOrStdout(), "No MCP servers found in Claude Code, Codex, Cursor, VS Code or Claude Desktop configs.") + return nil + } + selected, err := selectImports(cmd, discovered, args) + if err != nil { + return err + } + existing, _ := store.Servers(localmcp.ScopeUser) + taken := map[string]bool{} + for _, server := range existing { + taken[server.Name] = true + } + copied := 0 + states := env.State() + for _, server := range selected { + if taken[server.Name] { + _, _ = fmt.Fprintf(errOut, "skipped %s: already in your user config\n", server.Name) + continue + } + if _, err := localmcp.AddDisabledUserServer(store, states, server.Name, server.Config); err != nil { + _, _ = fmt.Fprintf(errOut, "skipped %s: %v\n", server.Name, err) + continue + } + taken[server.Name] = true + copied++ + _, _ = fmt.Fprintf(cmd.OutOrStdout(), "Imported %s from %s (off — enable it in /mcp)\n", server.Name, server.SourceLabel) + } + _, _ = fmt.Fprintf(cmd.OutOrStdout(), "Copied %d server(s) to %s\n", copied, scopeTarget(store, localmcp.ScopeUser)) + return nil + }, +} + +func selectImports(cmd *cobra.Command, discovered []localmcp.Server, names []string) ([]localmcp.Server, error) { + if importAll { + return discovered, nil + } + if len(names) > 0 { + wanted := map[string]bool{} + for _, name := range names { + wanted[name] = true + } + var selected []localmcp.Server + for _, server := range discovered { + if wanted[server.Name] { + selected = append(selected, server) + delete(wanted, server.Name) + } + } + if len(wanted) > 0 { + missing := make([]string, 0, len(wanted)) + for name := range wanted { + missing = append(missing, name) + } + sort.Strings(missing) + return nil, fmt.Errorf("not found in other tools' configs: %s", strings.Join(missing, ", ")) + } + return selected, nil + } + if !env.IsTTY() { + return nil, errors.New("pass server names or --all when not running in a terminal") + } + out := cmd.ErrOrStderr() + _, _ = fmt.Fprintln(out, "MCP servers found in other tools:") + for index, server := range discovered { + _, _ = fmt.Fprintf(out, " %2d. %-24s %s · from %s\n", index+1, server.Name, server.Config.Origin(), server.SourceLabel) + } + _, _ = fmt.Fprint(out, "Copy which into rootly? (e.g. 1,3 · all · Enter for none): ") + reader := bufio.NewReader(env.Stdin) + line, _ := reader.ReadString('\n') + return parseSelection(strings.TrimSpace(line), discovered) +} + +func parseSelection(input string, discovered []localmcp.Server) ([]localmcp.Server, error) { + if input == "" { + return nil, nil + } + if strings.EqualFold(input, "all") || input == "*" { + return discovered, nil + } + var selected []localmcp.Server + seen := map[int]bool{} + for _, field := range strings.FieldsFunc(input, func(r rune) bool { return r == ',' || r == ' ' }) { + number, err := strconv.Atoi(field) + if err != nil || number < 1 || number > len(discovered) { + return nil, fmt.Errorf("invalid selection %q: use numbers between 1 and %d", field, len(discovered)) + } + if !seen[number] { + seen[number] = true + selected = append(selected, discovered[number-1]) + } + } + return selected, nil +} + +func init() { + addCmd.Flags().StringVarP(&addTransport, "transport", "t", "", "Transport: stdio, http or sse (default stdio)") + addCmd.Flags().StringVarP(&addScope, "scope", "s", localmcp.ScopeLocal, "Config scope: local, user or project") + addCmd.Flags().StringArrayVarP(&addEnv, "env", "e", nil, "Environment variable for stdio servers (KEY=value, repeatable)") + addCmd.Flags().StringArrayVarP(&addHeaders, "header", "H", nil, `HTTP header ("Name: value", repeatable)`) + addCmd.Flags().StringVar(&addURL, "url", "", "Server URL (Codex syntax; implies --transport http)") + addCmd.Flags().StringVar(&addBearerEnv, "bearer-token-env-var", "", "Env var holding a bearer token for http servers (Codex syntax)") + addJSONCmd.Flags().StringVarP(&addJSONScope, "scope", "s", localmcp.ScopeLocal, "Config scope: local, user or project") + removeCmd.Flags().StringVarP(&removeScope, "scope", "s", "", "Config scope to remove from: local, user or project") + importCmd.Flags().BoolVar(&importAll, "all", false, "Copy every discovered server") + + MCPCmd.AddCommand(addCmd, addJSONCmd, listCmd, getCmd, removeCmd, loginCmd, logoutCmd, importCmd) +} diff --git a/internal/cmd/mcp/mcp_test.go b/internal/cmd/mcp/mcp_test.go new file mode 100644 index 0000000..a00663f --- /dev/null +++ b/internal/cmd/mcp/mcp_test.go @@ -0,0 +1,293 @@ +package mcp + +import ( + "bytes" + "encoding/json" + "errors" + "net/http" + "net/http/httptest" + "os" + "path/filepath" + "strings" + "testing" + + "github.com/spf13/pflag" + + localmcp "github.com/rootlyhq/rootly-cli/internal/mcp" +) + +type testEnv struct { + home string + project string +} + +func setupEnv(t *testing.T) testEnv { + t.Helper() + home, project := t.TempDir(), t.TempDir() + restore := SetEnvironment(Environment{ + CWD: func() (string, error) { return project, nil }, + Store: func(cwd string) *localmcp.Store { + return &localmcp.Store{UserPath: filepath.Join(home, ".rootly-cli", "mcp.json"), CWD: cwd} + }, + State: func() *localmcp.StateStore { + return &localmcp.StateStore{Path: filepath.Join(home, ".rootly-cli", "mcp-state.json")} + }, + Tokens: func() *localmcp.TokenStore { + return &localmcp.TokenStore{Path: filepath.Join(home, ".rootly-cli", "mcp-oauth.json")} + }, + Discovery: func(cwd string) localmcp.Discovery { return localmcp.Discovery{Home: home, CWD: cwd, GOOS: "linux"} }, + IsTTY: func() bool { return false }, + Stdin: strings.NewReader(""), + HTTP: http.DefaultClient, + OpenURL: func(string) error { return nil }, + }) + t.Cleanup(restore) + return testEnv{home: home, project: project} +} + +func resetFlags(flags *pflag.FlagSet) { + flags.VisitAll(func(flag *pflag.Flag) { + if slice, ok := flag.Value.(pflag.SliceValue); ok { + _ = slice.Replace(nil) + } else { + _ = flag.Value.Set(flag.DefValue) + } + flag.Changed = false + }) +} + +func run(t *testing.T, args ...string) (string, error) { + t.Helper() + for _, command := range MCPCmd.Commands() { + resetFlags(command.Flags()) + } + var out bytes.Buffer + MCPCmd.SetOut(&out) + MCPCmd.SetErr(&out) + MCPCmd.SetArgs(args) + err := MCPCmd.Execute() + return out.String(), err +} + +func readServers(t *testing.T, path string) map[string]any { + t.Helper() + data, err := os.ReadFile(path) + if err != nil { + t.Fatal(err) + } + var document map[string]any + if err := json.Unmarshal(data, &document); err != nil { + t.Fatal(err) + } + return document +} + +func TestAddScopesAndMaskedOutput(t *testing.T) { + env := setupEnv(t) + if out, err := run(t, "add", "-e", "TOKEN=secret-value", "fs", "--", "npx", "-y", "@modelcontextprotocol/server-filesystem", "/tmp"); err != nil { + t.Fatalf("%v: %s", err, out) + } + if out, err := run(t, "add", "-s", "project", "maestro", "--", "maestro", "mcp"); err != nil { + t.Fatalf("%v: %s", err, out) + } + if out, err := run(t, "add", "-s", "user", "-t", "sse", "-H", "Authorization: Bearer hidden", "events", "https://example.com/sse"); err != nil { + t.Fatalf("%v: %s", err, out) + } + if out, err := run(t, "add", "-s", "user", "--url", "https://mcp.linear.app/mcp", "--bearer-token-env-var", "LINEAR_TOKEN", "linear"); err != nil { + t.Fatalf("%v: %s", err, out) + } + + user := readServers(t, filepath.Join(env.home, ".rootly-cli", "mcp.json")) + local := user["projects"].(map[string]any)[env.project].(map[string]any)["mcpServers"].(map[string]any)["fs"].(map[string]any) + if local["command"] != "npx" || len(local["args"].([]any)) != 3 { + t.Fatalf("local fs = %v", local) + } + events := user["mcpServers"].(map[string]any)["events"].(map[string]any) + if events["type"] != "sse" || events["headers"].(map[string]any)["Authorization"] != "Bearer hidden" { + t.Fatalf("events = %v", events) + } + linear := user["mcpServers"].(map[string]any)["linear"].(map[string]any) + if linear["type"] != "http" || linear["bearer_token_env_var"] != "LINEAR_TOKEN" { + t.Fatalf("linear = %v", linear) + } + if readServers(t, filepath.Join(env.project, ".mcp.json"))["mcpServers"].(map[string]any)["maestro"] == nil { + t.Fatal("project server missing") + } + + list, err := run(t, "list") + if err != nil { + t.Fatal(err) + } + for _, want := range []string{"fs", "maestro", "events", "linear", "pending approval"} { + if !strings.Contains(list, want) { + t.Errorf("list missing %q:\n%s", want, list) + } + } + get, err := run(t, "get", "events") + if err != nil { + t.Fatal(err) + } + all := list + get + if fsGet, err := run(t, "get", "fs"); err == nil { + all += fsGet + } + if strings.Contains(all, "secret-value") || strings.Contains(all, "hidden") { + t.Fatalf("output leaks a secret:\n%s", all) + } + if !strings.Contains(all, localmcp.Mask) { + t.Fatalf("expected masked values:\n%s", all) + } + + if _, err := run(t, "remove", "fs"); err != nil { + t.Fatal(err) + } + if _, err := run(t, "get", "fs"); err == nil { + t.Fatal("removed server still found") + } +} + +func TestAddRejectsBadInput(t *testing.T) { + setupEnv(t) + cases := [][]string{ + {"add", "-H", "X: y", "fs", "--", "npx"}, + {"add", "-t", "ws", "x", "https://x"}, + {"add", "-s", "galaxy", "x", "--", "npx"}, + {"add", "-e", "NOEQUALS", "x", "--", "npx"}, + {"add", "bad name!", "--", "npx"}, + } + for _, args := range cases { + if _, err := run(t, args...); err == nil { + t.Errorf("%v should fail", args) + } + } +} + +func TestAddHTTPNeedsAuthentication(t *testing.T) { + setupEnv(t) + server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) { + w.Header().Set("WWW-Authenticate", `Bearer resource_metadata="https://example.com/.well-known/oauth-protected-resource"`) + w.WriteHeader(http.StatusUnauthorized) + })) + defer server.Close() + out, err := run(t, "add", "--transport", "http", "braintrust", server.URL+"/mcp") + if err != nil { + t.Fatalf("%v: %s", err, out) + } + if !strings.Contains(out, "needs authentication") || !strings.Contains(out, "rootly mcp login braintrust") { + t.Fatalf("out = %s", out) + } +} + +func TestAddJSONAndImport(t *testing.T) { + env := setupEnv(t) + if out, err := run(t, "add-json", "-s", "user", "everything", `{"command":"npx","args":["-y","@modelcontextprotocol/server-everything"]}`); err != nil { + t.Fatalf("%v: %s", err, out) + } + if _, err := run(t, "add-json", "broken", `{"type":"http"}`); err == nil { + t.Fatal("invalid json config should fail") + } + if err := os.MkdirAll(filepath.Join(env.project, ".vscode"), 0o755); err != nil { + t.Fatal(err) + } + if err := os.WriteFile(filepath.Join(env.project, ".vscode", "mcp.json"), []byte(`{"servers":{"github":{"type":"http","url":"https://api.githubcopilot.com/mcp/"}}}`), 0o600); err != nil { + t.Fatal(err) + } + if _, err := run(t, "import"); err == nil { + t.Fatal("import without a selection and no TTY should fail") + } + out, err := run(t, "import", "github") + if err != nil { + t.Fatalf("%v: %s", err, out) + } + servers := readServers(t, filepath.Join(env.home, ".rootly-cli", "mcp.json"))["mcpServers"].(map[string]any) + if servers["github"] == nil || servers["everything"] == nil { + t.Fatalf("servers = %v", servers) + } + if _, err := run(t, "remove", "-s", "imported", "github"); err == nil { + t.Fatal("imported servers are read-only") + } +} + +func TestImportStoresServerOffUntilEnabled(t *testing.T) { + testEnvironment := setupEnv(t) + if err := os.MkdirAll(filepath.Join(testEnvironment.project, ".vscode"), 0o755); err != nil { + t.Fatal(err) + } + config := `{"servers":{"github":{"command":"rootly-import-test-command","args":[]}}}` + if err := os.WriteFile(filepath.Join(testEnvironment.project, ".vscode", "mcp.json"), []byte(config), 0o600); err != nil { + t.Fatal(err) + } + out, err := run(t, "import", "github") + if err != nil { + t.Fatalf("%v: %s", err, out) + } + if !strings.Contains(out, "Imported github from ") || !strings.Contains(out, "(off — enable it in /mcp)") { + t.Fatalf("import output = %q, want disabled-state notice", out) + } + + store := env.Store(testEnvironment.project) + servers, err := store.Servers(localmcp.ScopeUser) + if err != nil { + t.Fatal(err) + } + if len(servers) != 1 || servers[0].Name != "github" || servers[0].Scope != localmcp.ScopeUser { + t.Fatalf("user-scope servers = %#v", servers) + } + state, err := env.State().Load() + if err != nil { + t.Fatal(err) + } + if state.IsEnabled(servers[0]) { + t.Fatal("imported user-scope server should be disabled") + } + + lookupStarted := make(chan struct{}, 1) + releaseLookup := make(chan struct{}) + manager := localmcp.NewManager(localmcp.ManagerOptions{ + CWD: testEnvironment.project, + Store: store, + State: env.State(), + Discovery: &localmcp.Discovery{Home: testEnvironment.home, CWD: testEnvironment.project, GOOS: "linux"}, + LookPath: func(name string) (string, error) { + if name == "rootly-import-test-command" { + lookupStarted <- struct{}{} + <-releaseLookup + return "", errors.New("blocked test command") + } + return "", errors.New("not installed") + }, + }) + defer manager.Close() + defer close(releaseLookup) + if err := manager.Load(); err != nil { + t.Fatal(err) + } + manager.Start() + if manager.Connecting() != 0 { + t.Fatal("manager should not start a disabled imported server") + } + for _, view := range manager.Views() { + if view.Server.ID() == servers[0].ID() && view.Status != localmcp.StatusOff { + t.Fatalf("imported server status = %q, want off", view.Status) + } + } + select { + case <-lookupStarted: + t.Fatal("manager tried to start the disabled imported server") + default: + } +} + +func TestListMasksSecretArgs(t *testing.T) { + setupEnv(t) + if out, err := run(t, "add", "-s", "user", "gh", "--", "gh-mcp", "--token=sk-live-123", "--api-key", "abc-456"); err != nil { + t.Fatalf("%v: %s", err, out) + } + list, err := run(t, "list") + if err != nil { + t.Fatal(err) + } + if strings.Contains(list, "sk-live-123") || strings.Contains(list, "abc-456") || !strings.Contains(list, "gh-mcp") { + t.Fatalf("list should mask secret args:\n%s", list) + } +} diff --git a/internal/cmd/mcp/search.go b/internal/cmd/mcp/search.go new file mode 100644 index 0000000..91ecd52 --- /dev/null +++ b/internal/cmd/mcp/search.go @@ -0,0 +1,83 @@ +package mcp + +import ( + "fmt" + "net/http" + "strings" + "time" + "unicode" + + "github.com/charmbracelet/x/ansi" + "github.com/spf13/cobra" + + localmcp "github.com/rootlyhq/rootly-cli/internal/mcp" +) + +var ( + registryURL = localmcp.RegistryURL + indexPath = localmcp.IndexPath + searchRefresh bool + searchLimit int +) + +var searchCmd = &cobra.Command{ + Use: "search [query]", + Short: "Search the MCP marketplace (official registry + curated list)", + Long: `Search the official MCP registry (registry.modelcontextprotocol.io) merged with a +curated list; Rootly-integrated vendors (Datadog, Sentry, Linear, Notion) come first. +The index is cached in ~/.rootly-cli/mcp-index.json for 7 days (--refresh refetches). +Nothing is installed: each result shows the exact command or URL and the +'rootly mcp add' line that adds it, or add it from /mcp in 'rootly agent'.`, + Args: cobra.ArbitraryArgs, + RunE: func(cmd *cobra.Command, args []string) error { + errOut := cmd.ErrOrStderr() + client := &http.Client{Timeout: 20 * time.Second} + catalog, err := localmcp.LoadCatalog(cmd.Context(), client, registryURL, indexPath(), searchRefresh, func(entries, pages int) { + _, _ = fmt.Fprintf(errOut, "\rIndexing MCP registry: %d servers (%d pages)", entries, pages) + }) + _, _ = fmt.Fprint(errOut, "\r\033[K") + if err != nil { + _, _ = fmt.Fprintf(errOut, "warning: %v\n", err) + } + results := localmcp.SearchMarket(catalog, strings.Join(args, " "), max(1, searchLimit)) + out := cmd.OutOrStdout() + if len(results) == 0 { + _, _ = fmt.Fprintln(out, "No MCP servers match.") + return nil + } + for _, entry := range results { + title := entry.Name + if entry.Author != "" { + title += " (" + entry.Author + ")" + } + _, _ = fmt.Fprintf(out, "%s\n %s\n runs: %s\n add: %s\n", oneLine(title), oneLine(entry.Description), plainText(localmcp.InstallLabel(entry)), plainText(localmcp.AddCommand(entry))) + if entry.EnvHint != "" { + _, _ = fmt.Fprintf(out, " needs: %s\n", oneLine(entry.EnvHint)) + } + } + return nil + }, +} + +func plainText(text string) string { + return strings.Map(func(r rune) rune { + if unicode.IsControl(r) || unicode.Is(unicode.Cf, r) { + return ' ' + } + return r + }, ansi.Strip(text)) +} + +func oneLine(text string) string { + text = strings.Join(strings.Fields(plainText(text)), " ") + if runes := []rune(text); len(runes) > 160 { + text = string(runes[:157]) + "..." + } + return text +} + +func init() { + searchCmd.Flags().BoolVar(&searchRefresh, "refresh", false, "Refetch the registry index") + searchCmd.Flags().IntVar(&searchLimit, "limit", 12, "Maximum results") + MCPCmd.AddCommand(searchCmd) +} diff --git a/internal/cmd/mcp/search_test.go b/internal/cmd/mcp/search_test.go new file mode 100644 index 0000000..5dee833 --- /dev/null +++ b/internal/cmd/mcp/search_test.go @@ -0,0 +1,76 @@ +package mcp + +import ( + "net/http" + "net/http/httptest" + "os" + "path/filepath" + "strings" + "testing" +) + +func TestSearchShowsExactInstallAndUserScopeAddLine(t *testing.T) { + test := setupEnv(t) + hits := 0 + server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) { + hits++ + _, _ = w.Write([]byte(`{"servers":[{"server":{"name":"io.github.acme/pager-mcp","description":"Pager tool","packages":[{"registryType":"npm","identifier":"@acme/pager-mcp","transport":{"type":"stdio"}}]},"_meta":{"io.modelcontextprotocol.registry/official":{"isLatest":true}}}],"metadata":{}}`)) + })) + defer server.Close() + previousURL, previousPath := registryURL, indexPath + cache := filepath.Join(test.home, ".rootly-cli", "mcp-index.json") + registryURL, indexPath = server.URL, func() string { return cache } + t.Cleanup(func() { registryURL, indexPath = previousURL, previousPath }) + + out, err := run(t, "search", "pager") + if err != nil { + t.Fatal(err) + } + for _, want := range []string{"pager-mcp (acme)", "runs: npx -y @acme/pager-mcp", "add: rootly mcp add -s user pager-mcp -- npx -y @acme/pager-mcp"} { + if !strings.Contains(out, want) { + t.Fatalf("missing %q in:\n%s", want, out) + } + } + if _, err := os.Stat(cache); err != nil { + t.Fatalf("index not cached: %v", err) + } + if _, err := run(t, "search", "sentry"); err != nil || hits != 1 { + t.Fatalf("second search should use the cache: hits=%d err=%v", hits, err) + } + out, _ = run(t, "search", "sentry") + if !strings.Contains(out, "runs: http https://mcp.sentry.dev/mcp") { + t.Fatalf("curated entry missing:\n%s", out) + } + if _, err := os.Stat(filepath.Join(test.home, ".rootly-cli", "mcp.json")); !os.IsNotExist(err) { + t.Fatal("search must not install anything") + } +} + +func TestSearchFallsBackToCuratedWhenRegistryFails(t *testing.T) { + test := setupEnv(t) + server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) { w.WriteHeader(http.StatusBadGateway) })) + defer server.Close() + previousURL, previousPath := registryURL, indexPath + registryURL, indexPath = server.URL, func() string { return filepath.Join(test.home, "idx.json") } + t.Cleanup(func() { registryURL, indexPath = previousURL, previousPath }) + out, err := run(t, "search", "datadog") + if err != nil { + t.Fatal(err) + } + if !strings.Contains(out, "warning: could not refresh the MCP registry index") || !strings.Contains(out, "datadog (datadog)") { + t.Fatalf("out:\n%s", out) + } +} + +func TestSearchStripsTerminalControls(t *testing.T) { + if got := oneLine("evil\x1b]0;pwned\x07 \x1b[2Jtext\u202e\r\nmore"); got != "evil text more" { + t.Fatalf("oneLine = %q", got) + } +} + +func TestSearchKeepsLongAddCommandsWhole(t *testing.T) { + args := strings.Repeat("--flag-value ", 20) + if got := plainText("rootly mcp add -s user x -- npx " + args); !strings.HasSuffix(got, args) { + t.Fatalf("add line truncated: %q", got) + } +} diff --git a/internal/cmd/mcp/select_test.go b/internal/cmd/mcp/select_test.go new file mode 100644 index 0000000..4e22e08 --- /dev/null +++ b/internal/cmd/mcp/select_test.go @@ -0,0 +1,160 @@ +package mcp + +import ( + "errors" + "os" + "path/filepath" + "strings" + "testing" + "time" + + localmcp "github.com/rootlyhq/rootly-cli/internal/mcp" +) + +func discoveredServers(names ...string) []localmcp.Server { + servers := make([]localmcp.Server, len(names)) + for index, name := range names { + servers[index] = localmcp.Server{Name: name, Config: localmcp.ServerConfig{Command: "npx"}, SourceLabel: "Claude Code"} + } + return servers +} + +func TestParseSelection(t *testing.T) { + discovered := discoveredServers("a", "b", "c") + cases := map[string][]string{ + "": nil, + "all": {"a", "b", "c"}, + "ALL": {"a", "b", "c"}, + "*": {"a", "b", "c"}, + "1,3": {"a", "c"}, + "3 1 3": {"c", "a"}, + " 2 , 2 ": {"b"}, + } + for input, want := range cases { + selected, err := parseSelection(strings.TrimSpace(input), discovered) + if err != nil { + t.Fatalf("parseSelection(%q): %v", input, err) + } + var got []string + for _, server := range selected { + got = append(got, server.Name) + } + if strings.Join(got, ",") != strings.Join(want, ",") { + t.Errorf("parseSelection(%q) = %v, want %v", input, got, want) + } + } + for _, input := range []string{"0", "4", "x", "1,-2"} { + if _, err := parseSelection(input, discovered); err == nil || !strings.Contains(err.Error(), "between 1 and 3") { + t.Errorf("parseSelection(%q) error = %v", input, err) + } + } +} + +func TestSelectImportsByNameAndWithoutTerminal(t *testing.T) { + discovered := discoveredServers("github", "sentry", "linear") + importAll = false + selected, err := selectImports(importCmd, discovered, []string{"linear", "github"}) + if err != nil || len(selected) != 2 || selected[0].Name != "github" || selected[1].Name != "linear" { + t.Fatalf("named selection = %+v, %v", selected, err) + } + if _, err := selectImports(importCmd, discovered, []string{"zeta", "github", "alpha"}); err == nil || + err.Error() != "not found in other tools' configs: alpha, zeta" { + t.Fatalf("missing names error = %v", err) + } + importAll = true + t.Cleanup(func() { importAll = false }) + if selected, err := selectImports(importCmd, discovered, nil); err != nil || len(selected) != 3 { + t.Fatalf("--all = %+v, %v", selected, err) + } + importAll = false + if _, err := selectImports(importCmd, discovered, nil); err == nil || !strings.Contains(err.Error(), "--all when not running in a terminal") { + t.Fatalf("non-terminal picker error = %v", err) + } +} + +func TestDescribeStateShowsScopeAndActivation(t *testing.T) { + setupEnv(t) + state, err := env.State().Load() + if err != nil { + t.Fatal(err) + } + tokens := env.Tokens() + imported := localmcp.Server{Name: "gh", Scope: localmcp.ScopeImported, SourceLabel: "Claude Code", Config: localmcp.ServerConfig{Command: "npx"}} + if got := describeState(state, tokens, imported); !strings.Contains(got, "imported from Claude Code, read-only") || !strings.Contains(got, "off") { + t.Errorf("imported server state = %q", got) + } + remote := localmcp.Server{Name: "remote", Scope: localmcp.ScopeUser, Config: localmcp.ServerConfig{Type: localmcp.TransportHTTP, URL: "https://mcp.example.com"}} + got := describeState(state, tokens, remote) + if got != "user" { + t.Errorf("remote user server state = %q", got) + } +} + +func TestAddRejectsConflictingTransportFlags(t *testing.T) { + setupEnv(t) + cases := []struct { + args []string + want string + }{ + {[]string{"add", "x", "--url", "https://mcp.example.com", "https://other.example.com"}, "--url takes the URL"}, + {[]string{"add", "x", "-H", "Authorization: Bearer y", "--", "npx", "server"}, "only apply to http and sse"}, + {[]string{"add", "x", "--transport", "sse", "a", "b"}, "sse servers take exactly one URL"}, + {[]string{"add", "x", "--transport", "http", "ftp://mcp.example.com"}, `invalid URL "ftp://mcp.example.com"`}, + {[]string{"add", "x", "-e", "TOKEN=1", "https://mcp.example.com"}, "-e/--env only applies to stdio"}, + {[]string{"add", "x", "--transport", "grpc", "y"}, `invalid transport "grpc"`}, + {[]string{"add", "x", "-e", "NOVALUE", "--", "npx"}, "KEY=value"}, + } + for _, tc := range cases { + if _, err := run(t, tc.args...); err == nil || !strings.Contains(err.Error(), tc.want) { + t.Errorf("rootly mcp %v error = %v, want %q", tc.args, err, tc.want) + } + } +} + +func TestAddHTTPWithCredentialsSkipsAuthenticationHint(t *testing.T) { + setupEnv(t) + out, err := run(t, "add", "x", "--url", "https://mcp.example.com", "-H", "Authorization: Bearer secret") + if err != nil { + t.Fatal(err) + } + if strings.Contains(out, "needs authentication") || strings.Contains(out, "secret") { + t.Fatalf("output = %q", out) + } + out, err = run(t, "add", "y", "--url", "https://mcp.example.com", "--bearer-token-env-var", "MCP_TOKEN") + if err != nil || strings.Contains(out, "needs authentication") { + t.Fatalf("bearer env var add = %q, %v", out, err) + } + if got := maskedValue("Authorization: Bearer secret", ":"); got != "Authorization:"+localmcp.Mask { + t.Errorf("maskedValue = %q", got) + } + if got := maskedValue("no-separator", ":"); got != "no-separator" { + t.Errorf("maskedValue without separator = %q", got) + } +} + +func TestCommandsFailWhenWorkingDirIsGone(t *testing.T) { + setupEnv(t) + previous := env.CWD + env.CWD = func() (string, error) { return "", errors.New("cwd removed") } + t.Cleanup(func() { env.CWD = previous }) + if _, err := run(t, "add", "x", "--", "npx", "server"); err == nil || !strings.Contains(err.Error(), "cwd removed") { + t.Fatalf("add with no cwd = %v", err) + } +} + +func TestDefaultEnvironmentUsesRealStores(t *testing.T) { + defaults := DefaultEnvironment() + if defaults.CWD == nil || defaults.Store == nil || defaults.State == nil || defaults.Tokens == nil || defaults.HTTP == nil || defaults.IsTTY == nil { + t.Fatalf("default environment has nil members: %+v", defaults) + } + if defaults.HTTP.Timeout != 5*time.Second { + t.Errorf("default HTTP timeout = %v", defaults.HTTP.Timeout) + } + home, _ := os.UserHomeDir() + if got := displayPath(filepath.Join(home, ".rootly-cli", "mcp.json")); got != filepath.Join("~", ".rootly-cli", "mcp.json") { + t.Errorf("displayPath(home file) = %q", got) + } + if got := displayPath("/srv/project/.mcp.json"); got != "/srv/project/.mcp.json" { + t.Errorf("displayPath(outside home) = %q", got) + } +} diff --git a/internal/cmd/mcp_register.go b/internal/cmd/mcp_register.go new file mode 100644 index 0000000..9db06f8 --- /dev/null +++ b/internal/cmd/mcp_register.go @@ -0,0 +1,7 @@ +package cmd + +import "github.com/rootlyhq/rootly-cli/internal/cmd/mcp" + +func init() { + rootCmd.AddCommand(mcp.MCPCmd) +} diff --git a/internal/cmd/root.go b/internal/cmd/root.go index 8e5ae14..4087486 100644 --- a/internal/cmd/root.go +++ b/internal/cmd/root.go @@ -1,6 +1,7 @@ package cmd import ( + "errors" "fmt" "os" @@ -133,12 +134,21 @@ func SetVersionInfo(v, c, d string) { version = v commit = c date = d + rootCmd.Version = v } // Execute runs the root command func Execute() { if err := rootCmd.Execute(); err != nil { fmt.Fprintf(os.Stderr, "Error: %v\n", err) - os.Exit(1) + os.Exit(exitCode(err)) } } + +func exitCode(err error) int { + var coder interface{ ExitCode() int } + if errors.As(err, &coder) && coder.ExitCode() > 0 { + return coder.ExitCode() + } + return 1 +} diff --git a/internal/mcp/builtin/builtin_test.go b/internal/mcp/builtin/builtin_test.go new file mode 100644 index 0000000..8039972 --- /dev/null +++ b/internal/mcp/builtin/builtin_test.go @@ -0,0 +1,554 @@ +package builtin + +import ( + "context" + "errors" + "os" + "os/exec" + "path/filepath" + "runtime" + "strings" + "testing" + "time" +) + +func TestFilesystemToolsConfinePathsAndDenySensitiveFiles(t *testing.T) { + home := t.TempDir() + t.Setenv("HOME", home) + root := filepath.Join(home, "project") + if err := os.MkdirAll(root, 0o755); err != nil { + t.Fatal(err) + } + outside := filepath.Join(home, "outside.txt") + if err := os.WriteFile(outside, []byte("outside"), 0o600); err != nil { + t.Fatal(err) + } + if err := os.Symlink(outside, filepath.Join(root, "outside-link")); err != nil { + t.Skipf("symlinks are unavailable: %v", err) + } + safeTarget := filepath.Join(root, "safe-config") + if err := os.WriteFile(safeTarget, []byte("not sensitive"), 0o600); err != nil { + t.Fatal(err) + } + if err := os.Symlink(safeTarget, filepath.Join(root, ".env.local")); err != nil { + t.Fatal(err) + } + registry := testRegistry(t, root, nil) + readFile := mustLookup(t, registry, "local_read_file") + for _, path := range []string{"../outside.txt", "outside-link", ".env.local"} { + result := readFile.Run(context.Background(), map[string]any{"path": path}) + if !result.IsError { + t.Errorf("reading %q unexpectedly succeeded: %q", path, result.Output) + } + } + listing := mustLookup(t, registry, "local_list_dir").Run(context.Background(), map[string]any{}) + if strings.Contains(listing.Output, "outside-link") || strings.Contains(listing.Output, ".env.local") { + t.Fatalf("directory listing returned a denied symlink: %s", listing.Output) + } + sensitiveRegistry := testRegistry(t, home, nil) + sensitiveReadFile := mustLookup(t, sensitiveRegistry, "local_read_file") + cases := []struct { + name string + path string + }{ + {name: "ssh directory", path: ".ssh/config"}, + {name: "aws directory", path: ".aws/config"}, + {name: "gcloud directory", path: ".config/gcloud/credentials.db"}, + {name: "kube config", path: ".kube/config"}, + {name: "env file", path: ".env"}, + {name: "env variant", path: ".env.production"}, + {name: "pem file", path: "server.pem"}, + {name: "key file", path: "server.key"}, + {name: "id file", path: "id_service"}, + } + for _, test := range cases { + t.Run(test.name, func(t *testing.T) { + fullPath := filepath.Join(home, test.path) + if err := os.MkdirAll(filepath.Dir(fullPath), 0o755); err != nil { + t.Fatal(err) + } + if err := os.WriteFile(fullPath, []byte("sensitive"), 0o600); err != nil { + t.Fatal(err) + } + result := sensitiveReadFile.Run(context.Background(), map[string]any{"path": test.path}) + if !result.IsError || !strings.Contains(result.Output, "path is denied") { + t.Errorf("sensitive path was not denied: %#v", result) + } + }) + } +} + +func TestSensitivePathComponentsAreDeniedForReadAndGlob(t *testing.T) { + root := t.TempDir() + for _, path := range []string{".env.private/creds.json", "keys.pem/x"} { + fullPath := filepath.Join(root, path) + if err := os.MkdirAll(filepath.Dir(fullPath), 0o755); err != nil { + t.Fatal(err) + } + if err := os.WriteFile(fullPath, []byte("sensitive"), 0o600); err != nil { + t.Fatal(err) + } + } + registry := testRegistry(t, root, nil) + readFile := mustLookup(t, registry, "local_read_file") + for _, path := range []string{".env.private/creds.json", "keys.pem/x"} { + result := readFile.Run(context.Background(), map[string]any{"path": path}) + if !result.IsError || !strings.Contains(result.Output, "path is denied") { + t.Errorf("read_file did not deny %q: %#v", path, result) + } + } + glob := mustLookup(t, registry, "local_glob") + result := glob.Run(context.Background(), map[string]any{"pattern": "**/*"}) + if strings.Contains(result.Output, ".env.private") || strings.Contains(result.Output, "keys.pem") { + t.Fatalf("glob returned children of denied directories: %q", result.Output) + } + for _, path := range []string{".env.private", "keys.pem"} { + result := glob.Run(context.Background(), map[string]any{"pattern": "**/*", "path": path}) + if !result.IsError || !strings.Contains(result.Output, "path is denied") { + t.Errorf("glob did not deny directory %q: %#v", path, result) + } + } +} + +func TestFilesystemTraversalSkipsGitAndSensitivePaths(t *testing.T) { + root := t.TempDir() + t.Setenv("HOME", root) + for path, body := range map[string]string{ + "visible.txt": "match", + ".git/hidden.txt": "match", + ".aws/credentials": "match", + "nested/visible-match.txt": "match", + } { + fullPath := filepath.Join(root, path) + if err := os.MkdirAll(filepath.Dir(fullPath), 0o755); err != nil { + t.Fatal(err) + } + if err := os.WriteFile(fullPath, []byte(body), 0o600); err != nil { + t.Fatal(err) + } + } + registry := testRegistry(t, root, nil) + glob := mustLookup(t, registry, "local_glob").Run(context.Background(), map[string]any{"pattern": "**/*.txt"}) + if strings.Contains(glob.Output, ".git") || strings.Contains(glob.Output, ".aws") { + t.Fatalf("glob returned a denied path: %s", glob.Output) + } + grep := mustLookup(t, registry, "local_grep").Run(context.Background(), map[string]any{ + "pattern": "match", + "output_mode": "content", + }) + if strings.Contains(grep.Output, ".git") || strings.Contains(grep.Output, ".aws") { + t.Fatalf("grep returned a denied path: %s", grep.Output) + } +} + +func TestFilesystemSearchUsesSymlinkPath(t *testing.T) { + root := t.TempDir() + if err := os.WriteFile(filepath.Join(root, "real.txt"), []byte("needle\n"), 0o600); err != nil { + t.Fatal(err) + } + if err := os.Symlink("real.txt", filepath.Join(root, "alias.txt")); err != nil { + t.Skipf("symlinks are unavailable: %v", err) + } + registry := testRegistry(t, root, nil) + glob := mustLookup(t, registry, "local_glob").Run(context.Background(), map[string]any{"pattern": "alias*"}) + if glob.Output != "alias.txt" { + t.Fatalf("glob output = %q, want alias.txt", glob.Output) + } + grep := mustLookup(t, registry, "local_grep").Run(context.Background(), map[string]any{ + "pattern": "needle", + "output_mode": "content", + "glob": "alias.txt", + }) + if grep.Output != "alias.txt:1:needle" { + t.Fatalf("grep output = %q, want alias.txt:1:needle", grep.Output) + } + listing := mustLookup(t, registry, "local_list_dir").Run(context.Background(), map[string]any{}) + if !strings.Contains(listing.Output, "alias.txt") { + t.Fatalf("directory listing omitted in-root symlink: %q", listing.Output) + } +} + +func TestOpenRootFileRejectsSymlinkToOutsideParent(t *testing.T) { + rootPath := t.TempDir() + outsidePath := t.TempDir() + if err := os.WriteFile(filepath.Join(outsidePath, "secret.txt"), []byte("outside"), 0o600); err != nil { + t.Fatal(err) + } + if err := os.Symlink(outsidePath, filepath.Join(rootPath, "parent")); err != nil { + t.Skipf("symlinks are unavailable: %v", err) + } + root, err := os.OpenRoot(rootPath) + if err != nil { + t.Fatal(err) + } + defer func() { _ = root.Close() }() + file, err := openRootFile(root, filepath.Join("parent", "secret.txt")) + if err == nil { + _ = file.Close() + t.Fatal("rooted open succeeded through a symlink to an outside parent") + } +} + +func TestGrepCapsMatchesWithinOneFile(t *testing.T) { + root := t.TempDir() + content := strings.Repeat("match\n", maxSearchItems+100) + if err := os.WriteFile(filepath.Join(root, "many.txt"), []byte(content), 0o600); err != nil { + t.Fatal(err) + } + result := mustLookup(t, testRegistry(t, root, nil), "local_grep").Run(context.Background(), map[string]any{ + "pattern": "match", + "output_mode": "content", + }) + matches := strings.Split(result.Output, "\n") + if len(matches) != maxSearchItems { + t.Fatalf("got %d matches, want %d", len(matches), maxSearchItems) + } + if !strings.HasSuffix(matches[maxSearchItems-1], ":200:match") || strings.Contains(result.Output, ":201:match") { + t.Fatalf("grep did not stop at the 200th match: %q", matches[maxSearchItems-1]) + } +} + +func TestFilesystemOutputIsCapped(t *testing.T) { + root := t.TempDir() + content := strings.Repeat("x", maxFSOutput+100) + if err := os.WriteFile(filepath.Join(root, "large.txt"), []byte(content), 0o600); err != nil { + t.Fatal(err) + } + result := mustLookup(t, testRegistry(t, root, nil), "local_read_file").Run( + context.Background(), + map[string]any{"path": "large.txt"}, + ) + if len([]rune(result.Output)) != maxFSOutput || !strings.Contains(result.Output, "[truncated]") { + t.Fatalf("expected capped output, got %d characters", len([]rune(result.Output))) + } +} + +func TestBashIsApprovalGatedAndEnforcesTimeout(t *testing.T) { + bash := mustLookup(t, testRegistry(t, t.TempDir(), nil), "local_bash") + if !bash.NeedsApproval(map[string]any{"command": "rm -f scratch"}) { + t.Fatal("bash must require approval for commands that change things") + } + result := bash.Run(context.Background(), map[string]any{ + "command": "exec sleep 5", + "timeout_seconds": 1, + }) + if !result.IsError || !strings.Contains(result.Output, "timed out") { + t.Fatalf("expected bash timeout, got %#v", result) + } + capped := bash.Run(context.Background(), map[string]any{"command": "printf '%30000s' ''"}) + if len([]rune(capped.Output)) != maxShellOutput || !strings.Contains(capped.Output, "[truncated]") { + t.Fatalf("expected bash output cap, got %d characters", len([]rune(capped.Output))) + } +} + +func TestShellTimeoutClampsSeconds(t *testing.T) { + cases := []struct { + name string + value any + want time.Duration + }{ + {name: "too large", value: int64(1 << 62), want: maxShellTimeout}, + {name: "below minimum", value: int64(0), want: time.Second}, + } + for _, test := range cases { + t.Run(test.name, func(t *testing.T) { + if got := shellTimeout(map[string]any{"timeout_seconds": test.value}); got != test.want { + t.Fatalf("shellTimeout() = %s, want %s", got, test.want) + } + }) + } +} + +func TestRunCommandWaitDelayForBackgroundedChild(t *testing.T) { + started := time.Now() + result := runCommand( + context.Background(), + t.TempDir(), + "bash", + []string{"-c", "sleep 30 & echo hi"}, + 10*time.Second, + maxShellOutput, + "bash", + ) + if elapsed := time.Since(started); elapsed > 5*time.Second { + t.Fatalf("command took %s with a backgrounded child", elapsed) + } + if result.IsError { + t.Fatalf("successful shell with a backgrounded child was marked as an error: %q", result.Output) + } + if !strings.Contains(result.Output, "(exit 0)") { + t.Fatalf("backgrounded command output did not report exit 0: %q", result.Output) + } + if !strings.Contains(result.Output, "hi") { + t.Fatalf("backgrounded command output = %q, want it to include hi", result.Output) + } + if !strings.Contains(result.Output, "[background process still running; output detached]") { + t.Fatalf("backgrounded command output omitted detached-output note: %q", result.Output) + } +} + +func TestBashRedactsSensitiveEnvironmentValues(t *testing.T) { + const secret = "supersecretvalue123" + t.Setenv("ROOTLY_TEST_TOKEN", secret) + bash := mustLookup(t, testRegistry(t, t.TempDir(), nil), "local_bash") + result := bash.Run(context.Background(), map[string]any{"command": `echo "$ROOTLY_TEST_TOKEN"`}) + if result.IsError { + t.Fatalf("bash command failed: %#v", result) + } + if strings.Contains(result.Output, secret) || !strings.Contains(result.Output, "[REDACTED]") { + t.Fatalf("sensitive environment value was not redacted: %q", result.Output) + } +} + +func TestBashRedactsShortSensitiveEnvironmentValues(t *testing.T) { + t.Setenv("DB_PASSWORD", "abc1") + t.Setenv("FEATURE_TOKEN_ENABLED", "true") + bash := mustLookup(t, testRegistry(t, t.TempDir(), nil), "local_bash") + result := bash.Run(context.Background(), map[string]any{ + "command": `printf '%s\n' "$DB_PASSWORD" "$FEATURE_TOKEN_ENABLED"`, + }) + if result.IsError { + t.Fatalf("bash command failed: %#v", result) + } + if strings.Contains(result.Output, "abc1") || !strings.Contains(result.Output, "[REDACTED]") { + t.Fatalf("short sensitive environment value was not redacted: %q", result.Output) + } + if !strings.Contains(result.Output, "true") { + t.Fatalf("common boolean value was redacted: %q", result.Output) + } +} + +func TestSensitiveEnvironmentValuesSkipsCommonFlags(t *testing.T) { + for _, value := range []string{"true", "FALSE", "yes", "NO", "none", "NULL"} { + t.Run(value, func(t *testing.T) { + secrets := sensitiveEnvironmentValues([]string{"FEATURE_TOKEN_ENABLED=" + value}) + if len(secrets) != 0 { + t.Fatalf("common flag value was marked sensitive: %#v", secrets) + } + }) + } +} + +func TestRunCommandRedactsBeforeOutputTruncation(t *testing.T) { + const secret = "supersecretvalue123" + t.Setenv("ROOTLY_TEST_TOKEN", secret) + result := runCommand( + context.Background(), + t.TempDir(), + "bash", + []string{"-c", `printf '%19975s' ''; printf '%s' "$ROOTLY_TEST_TOKEN"`}, + defaultTimeout, + maxShellOutput, + "bash", + ) + if strings.Contains(result.Output, secret[:5]) { + t.Fatalf("output truncation exposed part of a sensitive value: %q", result.Output) + } +} + +func TestCLIDiscoveryAllowListApprovalAndTimeout(t *testing.T) { + root := t.TempDir() + binDir := t.TempDir() + executable, argsExecutable, timeoutExecutable := createTestCLIs(t, binDir) + lookPath := func(bin string) (string, error) { + if bin == "gh" { + return executable, nil + } + return "", errors.New("not found") + } + registry := testRegistry(t, root, lookPath) + if _, ok := registry.Lookup("local_gh"); !ok { + t.Fatal("gh should be advertised when installed on PATH") + } + if _, ok := registry.Lookup("local_aws"); ok { + t.Fatal("aws should not be advertised when absent from PATH") + } + gh := mustLookup(t, registry, "local_gh") + result := gh.Run(context.Background(), map[string]any{"args": []any{"version"}}) + if result.IsError || len([]rune(result.Output)) != maxCLIOutput || !strings.Contains(result.Output, "[truncated]") { + t.Fatalf("expected capped CLI output, got %#v", result) + } + argsTool := newCLITool(root, cliSpec{bin: "gh"}, func(string) (string, error) { return argsExecutable, nil }, nil) + argument := "; touch " + filepath.Join(root, "injected") + argsResult := argsTool.Run(context.Background(), map[string]any{"args": []any{argument}}) + if argsResult.IsError || !strings.Contains(argsResult.Output, argument) { + t.Fatalf("CLI argument was not passed verbatim: %#v", argsResult) + } + if _, err := os.Stat(filepath.Join(root, "injected")); !os.IsNotExist(err) { + t.Fatalf("CLI argument was interpreted by a shell: %v", err) + } + timeoutTool := newCLITool(root, cliSpec{bin: "gh"}, func(string) (string, error) { return timeoutExecutable, nil }, nil) + ctx, cancel := context.WithTimeout(context.Background(), 30*time.Millisecond) + defer cancel() + timedOut := timeoutTool.Run(ctx, map[string]any{"args": []any{"version"}}) + if !timedOut.IsError || !strings.Contains(timedOut.Output, "timed out") { + t.Fatalf("expected context timeout, got %#v", timedOut) + } + notInstalled := newCLITool(root, cliSpec{bin: "gh"}, func(string) (string, error) { return "", errors.New("missing") }, nil) + if got := notInstalled.Run(context.Background(), map[string]any{"args": []any{}}); !got.IsError || !strings.Contains(got.Output, "not on PATH") { + t.Fatalf("expected missing CLI rejection, got %#v", got) + } + if got := runCLI(context.Background(), root, "sh", nil, lookPath); !got.IsError || !strings.Contains(got.Output, "not allow-listed") { + t.Fatalf("expected non-allow-listed CLI rejection, got %#v", got) + } +} + +func TestCLIApprovalHeuristics(t *testing.T) { + cases := []struct { + bin string + args []string + want bool + }{ + {bin: "aws", args: []string{"ec2", "describe-instances"}, want: false}, + {bin: "aws", args: []string{"--profile", "get-prod", "ec2", "terminate-instances"}, want: true}, + {bin: "aws", args: []string{"--region", "us-east-1", "ec2", "describe-instances"}, want: false}, + {bin: "aws", args: []string{"--weird", "x", "ec2", "describe-instances"}, want: true}, + {bin: "aws", args: []string{"--profile"}, want: true}, + {bin: "aws", args: []string{"ec2", "delete-vpc"}, want: true}, + {bin: "aws", args: []string{"secretsmanager", "get-secret-value"}, want: true}, + {bin: "aws", args: []string{"ssm", "get-parameter"}, want: true}, + {bin: "aws", args: []string{"ecr", "get-login-password"}, want: true}, + {bin: "aws", args: []string{"s3api", "get-object"}, want: true}, + {bin: "aws", args: []string{"kms", "decrypt"}, want: true}, + {bin: "aws", args: []string{"ec2", "describe-instances", "--with-decryption"}, want: true}, + {bin: "kubectl", args: []string{"get", "pods"}, want: false}, + {bin: "kubectl", args: []string{"get", "services"}, want: false}, + {bin: "kubectl", args: []string{"get", "svc"}, want: false}, + {bin: "kubectl", args: []string{"describe", "deploy/web"}, want: false}, + {bin: "kubectl", args: []string{"logs", "pod/x"}, want: false}, + {bin: "kubectl", args: []string{"get", "pods", "--kubeconfig"}, want: true}, + {bin: "kubectl", args: []string{"get", "pods", "--kubeconfig=/tmp/config"}, want: true}, + {bin: "kubectl", args: []string{"delete", "pod", "api"}, want: true}, + {bin: "kubectl", args: []string{"get", "secrets"}, want: true}, + {bin: "kubectl", args: []string{"get", "pods,secrets"}, want: true}, + {bin: "kubectl", args: []string{"describe", "secret/x"}, want: true}, + {bin: "kubectl", args: []string{"get", "-f", "x.yaml"}, want: true}, + {bin: "kubectl", args: []string{"get", "--filename=x.yaml"}, want: true}, + {bin: "kubectl", args: []string{"get", "-k", "overlays/prod"}, want: true}, + {bin: "kubectl", args: []string{"get", "--kustomize=overlays/prod"}, want: true}, + {bin: "kubectl", args: []string{"get", "--raw=/api"}, want: true}, + {bin: "gh", args: []string{"pr", "view", "42"}, want: false}, + {bin: "gh", args: []string{"pr", "view", "42", "--web"}, want: true}, + {bin: "gh", args: []string{"pr", "view", "42", "-w"}, want: true}, + {bin: "gh", args: []string{"issue", "list", "--web=1"}, want: true}, + {bin: "gh", args: []string{"pr", "merge", "42"}, want: true}, + {bin: "terraform", args: []string{"plan"}, want: true}, + {bin: "terraform", args: []string{"apply"}, want: true}, + {bin: "terraform", args: []string{"show"}, want: false}, + {bin: "terraform", args: []string{"show", "-json"}, want: false}, + {bin: "redis-cli", args: []string{"get", "app:key"}, want: false}, + {bin: "redis-cli", args: []string{"del", "app:key"}, want: true}, + {bin: "gcloud", args: []string{"compute", "instances", "list"}, want: true}, + {bin: "jq", args: []string{"."}, want: true}, + {bin: "pup", args: []string{"a"}, want: true}, + } + for _, test := range cases { + t.Run(test.bin+"/"+strings.Join(test.args, "_"), func(t *testing.T) { + if got := cliNeedsApproval(test.bin, test.args); got != test.want { + t.Fatalf("cliNeedsApproval(%q, %v) = %t, want %t", test.bin, test.args, got, test.want) + } + }) + } +} + +func TestKubectlKubeconfigInsideRootNeedsApproval(t *testing.T) { + root := t.TempDir() + configPath := filepath.Join(root, "config") + if err := os.WriteFile(configPath, []byte("{}"), 0o600); err != nil { + t.Fatal(err) + } + outsidePath := filepath.Join(t.TempDir(), "config") + if err := os.WriteFile(outsidePath, []byte("{}"), 0o600); err != nil { + t.Fatal(err) + } + t.Setenv("KUBECONFIG", outsidePath+string(os.PathListSeparator)+configPath) + kubectl := newCLITool(root, cliSpec{bin: "kubectl"}, func(string) (string, error) { return "", errors.New("unused") }, nil) + args := map[string]any{"args": []any{"get", "pods"}} + if !kubectl.NeedsApproval(args) { + t.Fatal("KUBECONFIG entry inside the working directory must require approval") + } + t.Setenv("KUBECONFIG", "") + if kubectl.NeedsApproval(args) { + t.Fatal("plain kubectl get pods should not require approval") + } +} + +func testRegistry(t *testing.T, root string, lookPath LookPathFunc) *Registry { + t.Helper() + if lookPath == nil { + lookPath = func(string) (string, error) { return "", errors.New("not installed") } + } + registry, err := NewRegistry(root, lookPath) + if err != nil { + t.Fatal(err) + } + return registry +} + +func mustLookup(t *testing.T, registry *Registry, name string) Tool { + t.Helper() + tool, ok := registry.Lookup(name) + if !ok { + t.Fatalf("tool %q was not registered", name) + } + return tool +} + +func createTestCLIs(t *testing.T, binDir string) (executable, argsExecutable, timeoutExecutable string) { + t.Helper() + sourcePath := filepath.Join(binDir, "fake_cli.go") + const source = `package main + +import ( + "fmt" + "os" + "path/filepath" + "strings" + "time" +) + +func main() { + name := strings.TrimSuffix(filepath.Base(os.Args[0]), filepath.Ext(os.Args[0])) + switch name { + case "args-gh": + if len(os.Args) > 1 { + fmt.Print(os.Args[1]) + } + case "slow-gh": + time.Sleep(5 * time.Second) + default: + fmt.Printf("%13000s", "") + } +} +` + if err := os.WriteFile(sourcePath, []byte(source), 0o600); err != nil { + t.Fatal(err) + } + workingDir, err := os.Getwd() + if err != nil { + t.Fatal(err) + } + extension := "" + if runtime.GOOS == "windows" { + extension = ".exe" + } + helperPath := filepath.Join(binDir, "cli-helper"+extension) + command := exec.CommandContext(t.Context(), "go", "build", "-o", helperPath, sourcePath) + command.Dir = workingDir + if output, err := command.CombinedOutput(); err != nil { + t.Fatalf("build fake CLI: %v\n%s", err, output) + } + helper, err := os.ReadFile(helperPath) + if err != nil { + t.Fatal(err) + } + paths := make([]string, 0, 3) + for _, name := range []string{"gh", "args-gh", "slow-gh"} { + path := filepath.Join(binDir, name+extension) + if err := os.WriteFile(path, helper, 0o700); err != nil { + t.Fatal(err) + } + paths = append(paths, path) + } + return paths[0], paths[1], paths[2] +} diff --git a/internal/mcp/builtin/cli.go b/internal/mcp/builtin/cli.go new file mode 100644 index 0000000..2964cb4 --- /dev/null +++ b/internal/mcp/builtin/cli.go @@ -0,0 +1,526 @@ +package builtin + +import ( + "bytes" + "context" + "encoding/json" + "errors" + "fmt" + "os" + "os/exec" + "path/filepath" + "regexp" + "sort" + "strings" + "time" + "unicode/utf8" +) + +const ( + maxCLIOutput = 12_000 + maxShellOutput = 20_000 + defaultTimeout = 30 * time.Second + maxShellTimeout = 600 * time.Second +) + +type cliSpec struct { + bin string + label string + description string +} + +type boundedBuffer struct { + buffer bytes.Buffer + limit int + truncated bool +} + +var cliRegistry = []cliSpec{ + {bin: "gh", label: "GitHub CLI", description: "Query and manage GitHub pull requests, issues, checks, runs, releases, and repositories."}, + {bin: "kubectl", label: "kubectl", description: "Inspect and manage Kubernetes resources, logs, events, and workloads."}, + {bin: "kubectx", label: "kubectx", description: "Show or switch Kubernetes contexts."}, + {bin: "aws", label: "AWS CLI", description: "Query and manage AWS resources using the configured local credentials."}, + {bin: "gcloud", label: "gcloud", description: "Query and manage Google Cloud resources using the configured local credentials."}, + {bin: "datadog", label: "Datadog CLI", description: "Query and manage Datadog monitors, metrics, events, and logs."}, + {bin: "dog", label: "Datadog dogshell", description: "Run Datadog dogshell commands."}, + {bin: "pup", label: "pup", description: "Parse and filter HTML from stdin using CSS selectors."}, + {bin: "jq", label: "jq", description: "Filter and transform JSON from stdin."}, + {bin: "psql", label: "psql", description: "Run SQL against the configured PostgreSQL server."}, + {bin: "redis-cli", label: "redis-cli", description: "Query Redis using the configured local credentials."}, + {bin: "terraform", label: "Terraform", description: "Inspect Terraform plans and state; mutating actions require approval."}, + {bin: "sentry-cli", label: "Sentry CLI", description: "Query and manage Sentry releases and issues."}, +} + +var sensitiveEnvironmentName = regexp.MustCompile(`(?i)(TOKEN|SECRET|PASSWORD|PASSWD|KEY|CREDENTIAL|AUTH|SESSION|COOKIE)`) + +func newCLITool(root string, spec cliSpec, lookPath LookPathFunc, policy *Policy) Tool { + name := "local_" + strings.ReplaceAll(spec.bin, "-", "_") + return localTool{ + name: name, + description: fmt.Sprintf("%s (%q) runs on this machine with its local credentials. Pass argument tokens via `args`.", spec.label, spec.bin), + schema: jsonSchema(`{"type":"object","properties":{"args":{"type":"array","items":{"type":"string"},"description":"CLI argument tokens, passed without a shell."}},"required":["args"]}`), + readOnly: false, + approve: func(args map[string]any) Approval { + values, ok := stringSliceArg(args, "args") + if !ok { + return Approval{Needs: true} + } + argv := append([]string{spec.bin}, values...) + builtinNeeds := cliNeedsApproval(spec.bin, values) || (spec.bin == "kubectl" && kubectlKubeconfigInsideRoot(root)) + return commandsApproval(policy, [][]string{argv}, func([]string) bool { return !builtinNeeds }) + }, + run: func(ctx context.Context, args map[string]any) Result { + values, ok := stringSliceArg(args, "args") + if !ok { + return errorResult("args must be an array of strings") + } + return runCLI(ctx, root, spec.bin, values, lookPath) + }, + } +} + +func newBashTool(root string, policy *Policy) Tool { + return localTool{ + name: "local_bash", + description: "Run a shell command on this machine with its local credentials. Read-only commands (ls, cat, grep, head, git status/log/diff/show) " + + "on files in the working directory run without approval; anything else asks first. Output and runtime are capped.", + schema: jsonSchema(`{"type":"object","properties":{"command":{"type":"string"},"timeout_seconds":{"type":"integer","minimum":1,"maximum":600}},"required":["command"]}`), + readOnly: false, + approve: func(args map[string]any) Approval { + command, _ := stringArg(args, "command") + if forbidden := scriptForbidden(policy, command); forbidden != "" { + return Approval{Forbidden: forbidden} + } + commands, ok := shellCommands(command) + if !ok { + return Approval{Needs: true} + } + return commandsApproval(policy, commands, func(argv []string) bool { return isKnownSafeCommand(argv, root) }) + }, + run: func(ctx context.Context, args map[string]any) Result { + command, ok := stringArg(args, "command") + if !ok || strings.TrimSpace(command) == "" { + return errorResult("command is required") + } + return runCommand(ctx, root, "bash", []string{"-c", hardenSafeGit(command, root)}, shellTimeout(args), maxShellOutput, "bash") + }, + } +} + +// commandsApproval combines user rules with the built-in read-only check: +// any forbidden command blocks the call; it runs unasked only when every +// command is allowed by a rule or known safe. +func commandsApproval(policy *Policy, commands [][]string, safe func([]string) bool) Approval { + approval := Approval{ReadOnly: true} + ruleAsks := false + for _, argv := range commands { + safeOnItsOwn := safe(argv) + approval.ReadOnly = approval.ReadOnly && safeOnItsOwn + decision, justification := policy.Evaluate(argv) + switch decision { + case DecisionForbidden: + if justification == "" { + justification = strings.Join(argv, " ") + } + return Approval{Forbidden: justification} + case DecisionPrompt: + ruleAsks = true + approval.Needs = true + case DecisionNone: + if !safeOnItsOwn { + approval.Needs = true + } + } + } + if approval.Needs && !ruleAsks && len(commands) == 1 { + approval.Prefix = rememberPrefix(commands[0]) + } + return approval +} + +// rememberPrefix is the command plus up to two leading subcommands, e.g. +// ["kubectl", "rollout", "undo"]. Commands that start with a flag or a path +// get no prefix, so they can only be approved once, as do wrappers and +// interpreters whose prefix would cover arbitrary commands. +func rememberPrefix(argv []string) []string { + if len(argv) < 2 || strings.Contains(argv[0], "/") || unrememberableCommands[argv[0]] { + return nil + } + prefix := []string{argv[0]} + for _, token := range argv[1:min(len(argv), 3)] { + if token == "" || strings.HasPrefix(token, "-") || strings.ContainsAny(token, "/=.:@ ") { + break + } + prefix = append(prefix, token) + } + if len(prefix) < 2 { + return nil + } + return prefix +} + +var unrememberableCommands = map[string]bool{ + "sudo": true, "su": true, "doas": true, "env": true, "xargs": true, "eval": true, "exec": true, "command": true, + "builtin": true, "nohup": true, "nice": true, "time": true, "timeout": true, "watch": true, "ssh": true, + "sh": true, "bash": true, "zsh": true, "dash": true, "ksh": true, "fish": true, "source": true, + "python": true, "python3": true, "node": true, "ruby": true, "perl": true, "php": true, "deno": true, "bun": true, + "npx": true, "uvx": true, "docker": true, +} + +func shellTimeout(args map[string]any) time.Duration { + maxSeconds := int64(maxShellTimeout / time.Second) + seconds := int64(defaultTimeout / time.Second) + switch value := args["timeout_seconds"].(type) { + case int: + seconds = int64(value) + case int64: + seconds = value + case float64: + switch { + case value < 1: + seconds = 1 + case value > float64(maxSeconds): + seconds = maxSeconds + default: + seconds = int64(value) + } + } + seconds = min(max(seconds, int64(1)), maxSeconds) + return time.Duration(seconds) * time.Second +} + +func runCLI(ctx context.Context, root, bin string, args []string, lookPath LookPathFunc) Result { + if !isAllowedCLI(bin) { + return errorResult("CLI not allow-listed: %s", bin) + } + if lookPath == nil { + lookPath = exec.LookPath + } + path, err := lookPath(bin) + if err != nil { + return errorResult("CLI not on PATH: %s", bin) + } + absolute, err := filepath.Abs(path) + if err != nil { + return errorResult("resolve CLI: %v", err) + } + return runCommand(ctx, root, absolute, args, defaultTimeout, maxCLIOutput, bin) +} + +func runCommand(ctx context.Context, root, command string, args []string, timeout time.Duration, outputLimit int, displayName string) Result { + commandContext, cancel := context.WithTimeout(ctx, timeout) + defer cancel() + cmd := exec.CommandContext(commandContext, command, args...) + cmd.WaitDelay = 2 * time.Second + cmd.Dir = root + environment := os.Environ() + cmd.Env = environment + secrets := sensitiveEnvironmentValues(environment) + bufferLimit := outputLimit * 4 + for _, secret := range secrets { + bufferLimit = max(bufferLimit, outputLimit*4+len(secret)) + } + stdout := &boundedBuffer{limit: bufferLimit} + stderr := &boundedBuffer{limit: bufferLimit} + cmd.Stdout = stdout + cmd.Stderr = stderr + err := cmd.Run() + exitCode := 0 + isError := err != nil + backgroundOutputDetached := errors.Is(err, exec.ErrWaitDelay) && commandContext.Err() == nil + if err != nil { + exitCode = 1 + var exitError *exec.ExitError + if errors.As(err, &exitError) { + exitCode = exitError.ExitCode() + } + } + if backgroundOutputDetached { + exitCode = cmd.ProcessState.ExitCode() + isError = exitCode != 0 + } + body := strings.TrimRight(stdout.String(), "\n") + if stderrText := strings.TrimRight(stderr.String(), "\n"); stderrText != "" { + if body != "" { + body += "\n" + } + body += "[stderr]\n" + stderrText + } + if backgroundOutputDetached { + if body != "" { + body += "\n" + } + body += "[background process still running; output detached]" + } + header := fmt.Sprintf("$ %s\n(exit %d)", displayName, exitCode) + if commandContext.Err() != nil { + header += "\n[command timed out]" + } + if stdout.truncated || stderr.truncated { + body += "\n[command output truncated]" + } + if body != "" { + header += "\n\n" + body + } + return Result{Output: capOutput(redactEnvironmentValues(header, secrets), outputLimit), IsError: isError} +} + +func sensitiveEnvironmentValues(environment []string) []string { + secrets := make([]string, 0) + for _, entry := range environment { + name, value, ok := strings.Cut(entry, "=") + if !ok || !sensitiveEnvironmentName.MatchString(name) || utf8.RuneCountInString(value) < 4 { + continue + } + switch strings.ToLower(value) { + case "true", "false", "yes", "no", "none", "null": + continue + } + secrets = append(secrets, value) + } + return secrets +} + +func redactEnvironmentValues(output string, secrets []string) string { + sort.Slice(secrets, func(i, j int) bool { return len(secrets[i]) > len(secrets[j]) }) + for _, secret := range secrets { + output = strings.ReplaceAll(output, secret, "[REDACTED]") + } + return output +} + +func (b *boundedBuffer) Write(data []byte) (int, error) { + remaining := b.limit - b.buffer.Len() + if remaining > 0 { + _, _ = b.buffer.Write(data[:min(len(data), remaining)]) + } + if len(data) > remaining { + b.truncated = true + } + return len(data), nil +} + +func (b *boundedBuffer) String() string { + return b.buffer.String() +} + +func isAllowedCLI(bin string) bool { + for _, spec := range cliRegistry { + if spec.bin == bin { + return true + } + } + return false +} + +func cliNeedsApproval(bin string, args []string) bool { + lower := make([]string, len(args)) + for index := range args { + lower[index] = strings.ToLower(args[index]) + } + switch bin { + case "gh": + return ghNeedsApproval(lower) + case "kubectl": + return kubectlNeedsApproval(lower) + case "aws": + return awsNeedsApproval(lower) + case "terraform": + return len(lower) == 0 || lower[0] != "show" + case "kubectx": + return len(lower) > 0 + case "jq", "pup": + return true + case "redis-cli": + return len(lower) == 0 || !containsCommand(lower[0], "get", "mget", "scan", "type", "exists", "ttl", "pttl", "info", "dbsize") + case "sentry-cli": + return len(lower) < 2 || !containsCommand(lower[0], "issues", "releases") || + !containsCommand(lower[1], "list", "info", "show") + default: + return true + } +} + +func ghNeedsApproval(args []string) bool { + if len(args) < 2 || !containsCommand(args[0], "pr", "issue") || !containsCommand(args[1], "view", "list") { + return true + } + for _, arg := range args { + if arg == "--web" || arg == "-w" || strings.HasPrefix(arg, "--web=") { + return true + } + } + return false +} + +func awsNeedsApproval(args []string) bool { + _, action, ok := awsServiceOperation(args) + if !ok { + return true + } + for _, sensitive := range []string{"secret", "password", "token", "credential", "login", "key", "decrypt"} { + if strings.Contains(action, sensitive) { + return true + } + } + if containsCommand(action, "get-object", "get-parameter", "get-parameters", "get-parameters-by-path") { + return true + } + for _, arg := range args { + if arg == "--with-decryption" { + return true + } + } + return !strings.HasPrefix(action, "describe-") && + !strings.HasPrefix(action, "list-") && + !strings.HasPrefix(action, "get-") +} + +func awsServiceOperation(args []string) (service, operation string, ok bool) { + positionals := make([]string, 0, 2) + for index := 0; index < len(args) && len(positionals) < 2; index++ { + arg := args[index] + if !strings.HasPrefix(arg, "-") { + positionals = append(positionals, arg) + continue + } + option, _, hasValue := strings.Cut(arg, "=") + if hasValue && awsValueOption(option) { + continue + } + if awsBooleanOption(arg) { + continue + } + if awsValueOption(arg) { + if index+1 >= len(args) { + return "", "", false + } + index++ + continue + } + return "", "", false + } + if len(positionals) < 2 { + return "", "", false + } + return positionals[0], positionals[1], true +} + +func awsValueOption(option string) bool { + switch option { + case "--profile", "--region", "--output", "--endpoint-url", "--query", + "--cli-read-timeout", "--cli-connect-timeout", "--ca-bundle", "--color", + "--cli-binary-format": + return true + default: + return false + } +} + +func awsBooleanOption(option string) bool { + switch option { + case "--debug", "--no-verify-ssl", "--no-paginate", "--no-sign-request", "--no-cli-pager": + return true + default: + return false + } +} + +func kubectlNeedsApproval(args []string) bool { + for _, arg := range args { + if arg == "--kubeconfig" || strings.HasPrefix(arg, "--kubeconfig=") { + return true + } + } + if len(args) == 0 || !containsCommand(args[0], "get", "describe", "logs") { + return true + } + for _, arg := range args { + if strings.Contains(arg, "secret") || + strings.HasPrefix(arg, "-f") || + strings.HasPrefix(arg, "--filename") || + strings.HasPrefix(arg, "-k") || + strings.HasPrefix(arg, "--kustomize") || + strings.HasPrefix(arg, "--raw") { + return true + } + if strings.HasPrefix(arg, "-") { + continue + } + for _, part := range strings.FieldsFunc(strings.ToLower(arg), func(r rune) bool { + return r == ',' || r == '/' + }) { + if part == "se" || part == "secret" || part == "secrets" || + strings.HasSuffix(part, ".secrets") || + strings.HasPrefix(part, "secrets.") || + strings.HasPrefix(part, "secret.") { + return true + } + } + } + return false +} + +func kubectlKubeconfigInsideRoot(root string) bool { + rootPath, err := filepath.Abs(root) + if err != nil { + return false + } + rootPath, err = filepath.EvalSymlinks(rootPath) + if err != nil { + return false + } + for _, configPath := range filepath.SplitList(os.Getenv("KUBECONFIG")) { + if configPath == "" { + continue + } + if !filepath.IsAbs(configPath) { + configPath = filepath.Join(rootPath, configPath) + } + configPath, err = filepath.Abs(configPath) + if err != nil { + continue + } + configPath, err = filepath.EvalSymlinks(configPath) + if err == nil && isWithin(rootPath, configPath) { + return true + } + } + return false +} + +func containsCommand(value string, allowed ...string) bool { + for _, candidate := range allowed { + if value == candidate { + return true + } + } + return false +} + +func stringSliceArg(args map[string]any, key string) ([]string, bool) { + value, ok := args[key] + if !ok { + return nil, false + } + switch values := value.(type) { + case []string: + return append([]string(nil), values...), true + case []any: + result := make([]string, len(values)) + for index, value := range values { + text, ok := value.(string) + if !ok || strings.ContainsRune(text, '\x00') { + return nil, false + } + result[index] = text + } + return result, true + default: + return nil, false + } +} + +func jsonSchema(value string) json.RawMessage { + return json.RawMessage(value) +} diff --git a/internal/mcp/builtin/cli_kubectl_test.go b/internal/mcp/builtin/cli_kubectl_test.go new file mode 100644 index 0000000..cb3de50 --- /dev/null +++ b/internal/mcp/builtin/cli_kubectl_test.go @@ -0,0 +1,33 @@ +package builtin + +import ( + "testing" +) + +func TestKubectlSecretAliasApprovalHeuristics(t *testing.T) { + cases := []struct { + name string + args []string + want bool + }{ + {name: "get pods", args: []string{"get", "pods"}, want: false}, + {name: "get services", args: []string{"get", "services"}, want: false}, + {name: "get svc", args: []string{"get", "svc"}, want: false}, + {name: "describe deployment", args: []string{"describe", "deploy/web"}, want: false}, + {name: "pod logs", args: []string{"logs", "pod/x"}, want: false}, + {name: "short secret alias", args: []string{"get", "se", "x", "-o", "yaml"}, want: true}, + {name: "uppercase short alias", args: []string{"get", "SE"}, want: true}, + {name: "comma-separated alias", args: []string{"get", "pods,se"}, want: true}, + {name: "slash-separated alias", args: []string{"get", "se/foo"}, want: true}, + {name: "plural group-qualified alias", args: []string{"describe", "secrets.v1", "x"}, want: true}, + {name: "singular group-qualified alias", args: []string{"describe", "secret.v1", "x"}, want: true}, + {name: "case-insensitive Secret", args: []string{"get", "Secret"}, want: true}, + } + for _, test := range cases { + t.Run(test.name, func(t *testing.T) { + if got := cliNeedsApproval("kubectl", test.args); got != test.want { + t.Fatalf("cliNeedsApproval(kubectl, %v) = %t, want %t", test.args, got, test.want) + } + }) + } +} diff --git a/internal/mcp/builtin/fs.go b/internal/mcp/builtin/fs.go new file mode 100644 index 0000000..a2e6049 --- /dev/null +++ b/internal/mcp/builtin/fs.go @@ -0,0 +1,580 @@ +package builtin + +import ( + "bufio" + "context" + "errors" + "fmt" + "io" + "io/fs" + "os" + "path" + "path/filepath" + "regexp" + "sort" + "strings" +) + +const ( + maxFileBytes = 256 * 1024 + maxFileLines = 2000 + maxSearchItems = 200 + maxDirEntries = 500 + maxFSOutput = 20_000 +) + +var errSearchLimit = errors.New("search result limit reached") + +func newReadFileTool(root string) Tool { + return localTool{ + name: "local_read_file", + description: "Read a file inside the current working directory. Returns numbered lines.", + schema: jsonSchema(`{"type":"object","properties":{"path":{"type":"string"},"offset":{"type":"integer","minimum":0},"limit":{"type":"integer","minimum":1}},"required":["path"]}`), + readOnly: true, + needsApproval: func(map[string]any) bool { return false }, + run: func(_ context.Context, args map[string]any) Result { + filePath, ok := stringArg(args, "path") + if !ok || strings.TrimSpace(filePath) == "" { + return errorResult("path is required") + } + resolved, err := resolvePath(root, filePath) + if err != nil { + return errorResult("%v", err) + } + if deniedPath(resolved) { + return errorResult("path is denied") + } + rootFS, err := os.OpenRoot(root) + if err != nil { + return errorResult("%v", err) + } + defer func() { _ = rootFS.Close() }() + relative, err := rootRelativePath(root, resolved) + if err != nil { + return errorResult("%v", err) + } + info, err := rootFS.Stat(relative) + if err != nil { + return errorResult("%v", err) + } + if !info.Mode().IsRegular() { + return errorResult("path is not a regular file") + } + file, err := openRootFile(rootFS, relative) + if err != nil { + return errorResult("%v", err) + } + fileInfo, statErr := file.Stat() + if statErr != nil { + _ = file.Close() + return errorResult("%v", statErr) + } + if !fileInfo.Mode().IsRegular() { + _ = file.Close() + return errorResult("path is not a regular file") + } + data, readErr := io.ReadAll(io.LimitReader(file, maxFileBytes+1)) + closeErr := file.Close() + if readErr != nil { + return errorResult("%v", readErr) + } + if closeErr != nil { + return errorResult("%v", closeErr) + } + truncated := len(data) > maxFileBytes + if truncated { + data = data[:maxFileBytes] + } + lines := strings.Split(string(data), "\n") + offset := intArg(args, "offset", 0) + limit := intArg(args, "limit", maxFileLines) + if offset < 0 || limit < 1 { + return errorResult("offset must be non-negative and limit must be positive") + } + if limit > maxFileLines { + limit = maxFileLines + } + start := min(offset, len(lines)) + end := min(start+limit, len(lines)) + var output strings.Builder + for index := start; index < end; index++ { + fmt.Fprintf(&output, "%6d\t%s", index+1, lines[index]) + if index+1 < end { + output.WriteByte('\n') + } + } + if end < len(lines) || truncated { + fmt.Fprintf(&output, "\n…[%d more lines]", max(0, len(lines)-end)) + } + return Result{Output: capOutput(output.String(), maxFSOutput)} + }, + } +} + +func newGrepTool(root string) Tool { + return localTool{ + name: "local_grep", + description: "Search files inside the current working directory. Supports files_with_matches, content, or count output.", + schema: jsonSchema(`{"type":"object","properties":{"pattern":{"type":"string"},"path":{"type":"string"},"glob":{"type":"string"},"output_mode":{"type":"string","enum":["files_with_matches","content","count"]}},"required":["pattern"]}`), + readOnly: true, + needsApproval: func(map[string]any) bool { return false }, + run: func(ctx context.Context, args map[string]any) Result { + pattern, ok := stringArg(args, "pattern") + if !ok || pattern == "" { + return errorResult("pattern is required") + } + expression, err := regexp.Compile(pattern) + if err != nil { + return errorResult("invalid pattern: %v", err) + } + searchPath, ok := stringArg(args, "path") + if !ok || strings.TrimSpace(searchPath) == "" { + searchPath = "." + } + start, err := resolvePath(root, searchPath) + if err != nil { + return errorResult("%v", err) + } + if deniedPath(start) { + return errorResult("path is denied") + } + rootFS, err := os.OpenRoot(root) + if err != nil { + return errorResult("%v", err) + } + defer func() { _ = rootFS.Close() }() + startRelative, err := rootRelativePath(root, start) + if err != nil { + return errorResult("%v", err) + } + info, err := rootFS.Stat(startRelative) + if err != nil { + return errorResult("%v", err) + } + if !info.IsDir() && !info.Mode().IsRegular() { + return errorResult("path is not a file or directory") + } + mode, ok := stringArg(args, "output_mode") + if !ok || mode == "" { + mode = "files_with_matches" + } + if mode != "files_with_matches" && mode != "content" && mode != "count" { + return errorResult("output_mode must be files_with_matches, content, or count") + } + glob, _ := stringArg(args, "glob") + if glob != "" && filepath.IsAbs(glob) { + return errorResult("glob must be relative") + } + results := make([]string, 0, maxSearchItems) + err = walkFiles(ctx, root, rootFS, start, func(filePath, relative string) error { + if glob != "" && !globMatch(glob, relativeTo(start, filePath)) { + return nil + } + if len(results) >= maxSearchItems { + return errSearchLimit + } + remaining := maxSearchItems - len(results) + matches, scanErr := grepFile(rootFS, relative, expression, mode, relativeTo(root, filePath), remaining) + if scanErr != nil { + return nil + } + if len(matches) == 0 { + return nil + } + for _, match := range matches { + if len(results) >= maxSearchItems { + return errSearchLimit + } + results = append(results, match) + if len(results) >= maxSearchItems { + return errSearchLimit + } + } + return nil + }) + if err != nil && !errors.Is(err, errSearchLimit) { + return errorResult("%v", err) + } + if len(results) == 0 { + return Result{Output: "(no matches)"} + } + return Result{Output: capOutput(strings.Join(results, "\n"), maxFSOutput)} + }, + } +} + +func newGlobTool(root string) Tool { + return localTool{ + name: "local_glob", + description: "Find files by glob pattern inside the current working directory. Returns at most 200 paths.", + schema: jsonSchema(`{"type":"object","properties":{"pattern":{"type":"string"},"path":{"type":"string"}},"required":["pattern"]}`), + readOnly: true, + needsApproval: func(map[string]any) bool { return false }, + run: func(ctx context.Context, args map[string]any) Result { + pattern, ok := stringArg(args, "pattern") + if !ok || strings.TrimSpace(pattern) == "" { + return errorResult("pattern is required") + } + if filepath.IsAbs(pattern) || escapesRoot(pattern) { + return errorResult("glob must stay inside the current working directory") + } + base, ok := stringArg(args, "path") + if !ok || strings.TrimSpace(base) == "" { + base = "." + } + start, err := resolvePath(root, base) + if err != nil { + return errorResult("%v", err) + } + if deniedPath(start) { + return errorResult("path is denied") + } + rootFS, err := os.OpenRoot(root) + if err != nil { + return errorResult("%v", err) + } + defer func() { _ = rootFS.Close() }() + startRelative, err := rootRelativePath(root, start) + if err != nil { + return errorResult("%v", err) + } + info, err := rootFS.Stat(startRelative) + if err != nil { + return errorResult("%v", err) + } + if !info.IsDir() { + return errorResult("path is not a directory") + } + matches := make([]string, 0, maxSearchItems) + err = walkFiles(ctx, root, rootFS, start, func(filePath, _ string) error { + if globMatch(filepath.ToSlash(pattern), filepath.ToSlash(relativeTo(start, filePath))) { + matches = append(matches, filepath.ToSlash(relativeTo(start, filePath))) + if len(matches) >= maxSearchItems { + return errSearchLimit + } + } + return nil + }) + if err != nil && !errors.Is(err, errSearchLimit) { + return errorResult("%v", err) + } + if len(matches) == 0 { + return Result{Output: "(no matches)"} + } + sort.Strings(matches) + return Result{Output: capOutput(strings.Join(matches, "\n"), maxFSOutput)} + }, + } +} + +func newListDirTool(root string) Tool { + return localTool{ + name: "local_list_dir", + description: "List the non-sensitive entries in a directory inside the current working directory.", + schema: jsonSchema(`{"type":"object","properties":{"path":{"type":"string"}}}`), + readOnly: true, + needsApproval: func(map[string]any) bool { return false }, + run: func(_ context.Context, args map[string]any) Result { + dir, ok := stringArg(args, "path") + if !ok || strings.TrimSpace(dir) == "" { + dir = "." + } + resolved, err := resolvePath(root, dir) + if err != nil { + return errorResult("%v", err) + } + if deniedPath(resolved) { + return errorResult("path is denied") + } + rootFS, err := os.OpenRoot(root) + if err != nil { + return errorResult("%v", err) + } + defer func() { _ = rootFS.Close() }() + relative, err := rootRelativePath(root, resolved) + if err != nil { + return errorResult("%v", err) + } + directoryFile, err := rootFS.Open(relative) + if err != nil { + return errorResult("%v", err) + } + entries, readErr := directoryFile.ReadDir(-1) + closeErr := directoryFile.Close() + if readErr != nil { + return errorResult("%v", readErr) + } + if closeErr != nil { + return errorResult("%v", closeErr) + } + sort.Slice(entries, func(i, j int) bool { return entries[i].Name() < entries[j].Name() }) + output := make([]string, 0, min(len(entries), maxDirEntries)) + for _, entry := range entries { + entryPath := filepath.Join(resolved, entry.Name()) + entryRelative := filepath.Join(relative, entry.Name()) + if _, statErr := rootFS.Stat(entryRelative); statErr != nil { + continue + } + target, targetErr := filepath.EvalSymlinks(entryPath) + if entry.Name() == ".git" || deniedPath(entryPath) || targetErr != nil || !isWithin(root, target) || deniedPath(target) { + continue + } + if len(output) == maxDirEntries { + break + } + name := entry.Name() + if entry.IsDir() { + name += "/" + } + output = append(output, name) + } + if len(output) == 0 { + return Result{Output: "(empty)"} + } + return Result{Output: capOutput(strings.Join(output, "\n"), maxFSOutput)} + }, + } +} + +func resolvePath(root, requested string) (string, error) { + if strings.TrimSpace(requested) == "" { + return "", errors.New("path is required") + } + root = resolveRoot(root) + requestedPath := requested + if !filepath.IsAbs(requestedPath) { + requestedPath = filepath.Join(root, requestedPath) + } + absolute, err := filepath.Abs(filepath.Clean(requestedPath)) + if err != nil { + return "", err + } + if !isWithin(root, absolute) { + return "", errors.New("path escapes the current working directory") + } + if deniedPath(absolute) { + return "", errors.New("path is denied") + } + resolved, err := filepath.EvalSymlinks(absolute) + if err != nil { + return "", fmt.Errorf("resolve path: %w", err) + } + if !isWithin(root, resolved) { + return "", errors.New("path escapes the current working directory") + } + return resolved, nil +} + +func rootRelativePath(root, validatedPath string) (string, error) { + relative, err := filepath.Rel(root, validatedPath) + if err != nil { + return "", err + } + if filepath.IsAbs(relative) || escapesRoot(relative) { + return "", errors.New("path escapes the current working directory") + } + return filepath.Clean(relative), nil +} + +func openRootFile(root *os.Root, relative string) (*os.File, error) { + return root.Open(relative) +} + +func isWithin(root, candidate string) bool { + relative, err := filepath.Rel(root, candidate) + return err == nil && relative != ".." && !strings.HasPrefix(relative, ".."+string(os.PathSeparator)) && !filepath.IsAbs(relative) +} + +func deniedPath(candidate string) bool { + clean := filepath.Clean(candidate) + userHome, _ := os.UserHomeDir() + for _, home := range []string{os.Getenv("HOME"), userHome} { + if home == "" { + continue + } + if absolute, err := filepath.Abs(home); err == nil { + home = absolute + } + protectedHomes := []string{home} + if resolved, err := filepath.EvalSymlinks(home); err == nil && resolved != home { + protectedHomes = append(protectedHomes, resolved) + } + for _, protectedHome := range protectedHomes { + for _, protected := range []string{ + filepath.Join(protectedHome, ".ssh"), + filepath.Join(protectedHome, ".aws"), + filepath.Join(protectedHome, ".config", "gcloud"), + filepath.Join(protectedHome, ".kube", "config"), + filepath.Join(protectedHome, ".rootly-cli"), + } { + if isWithin(protected, clean) { + return true + } + } + } + } + for parent := clean; ; parent = filepath.Dir(parent) { + name := filepath.Base(parent) + base := strings.ToLower(name) + if name == ".git" || strings.HasPrefix(base, ".env") || strings.HasSuffix(base, ".pem") || + strings.HasSuffix(base, ".key") || strings.HasPrefix(base, "id_") { + return true + } + if parent == filepath.Dir(parent) { + break + } + } + return false +} + +func walkFiles(ctx context.Context, rootPath string, root *os.Root, start string, visit func(string, string) error) error { + startRelative, err := rootRelativePath(rootPath, start) + if err != nil { + return err + } + walkStart := filepath.ToSlash(startRelative) + return fs.WalkDir(root.FS(), walkStart, func(walkPath string, entry fs.DirEntry, walkErr error) error { + if walkErr != nil { + return nil + } + if err := ctx.Err(); err != nil { + return err + } + filePath := filepath.Join(rootPath, filepath.FromSlash(walkPath)) + if walkPath != walkStart && deniedPath(filePath) { + if entry.IsDir() { + return fs.SkipDir + } + return nil + } + resolved, err := filepath.EvalSymlinks(filePath) + if err != nil || !isWithin(rootPath, resolved) || deniedPath(resolved) { + if entry.IsDir() { + return fs.SkipDir + } + return nil + } + relative, err := rootRelativePath(rootPath, filePath) + if err != nil { + return nil + } + info, err := root.Stat(relative) + if err != nil || !info.Mode().IsRegular() { + return nil + } + return visit(filePath, relative) + }) +} + +func grepFile(root *os.Root, relative string, expression *regexp.Regexp, mode, displayPath string, maxMatches int) ([]string, error) { + file, err := openRootFile(root, relative) + if err != nil { + return nil, err + } + defer func() { _ = file.Close() }() + scanner := bufio.NewScanner(io.LimitReader(file, maxFileBytes)) + scanner.Buffer(make([]byte, 64*1024), maxFileBytes) + var matches []string + count := 0 + lineNumber := 0 + for scanner.Scan() { + lineNumber++ + line := scanner.Text() + if !expression.MatchString(line) { + continue + } + count++ + switch mode { + case "content": + matches = append(matches, fmt.Sprintf("%s:%d:%s", displayPath, lineNumber, line)) + if len(matches) >= maxMatches { + return matches, nil + } + case "files_with_matches": + return []string{displayPath}, nil + } + } + if err := scanner.Err(); err != nil { + return nil, err + } + if mode == "count" && count > 0 { + matches = append(matches, fmt.Sprintf("%s:%d", displayPath, count)) + } + return matches, nil +} + +func globMatch(pattern, value string) bool { + patternParts := strings.Split(filepath.ToSlash(pattern), "/") + valueParts := strings.Split(filepath.ToSlash(value), "/") + return globSegmentsMatch(patternParts, valueParts) +} + +func globSegmentsMatch(pattern, value []string) bool { + if len(pattern) == 0 { + return len(value) == 0 + } + if pattern[0] == "**" { + for index := 0; index <= len(value); index++ { + if globSegmentsMatch(pattern[1:], value[index:]) { + return true + } + } + return false + } + if len(value) == 0 { + return false + } + matched, err := path.Match(pattern[0], value[0]) + return err == nil && matched && globSegmentsMatch(pattern[1:], value[1:]) +} + +func relativeTo(base, value string) string { + relative, err := filepath.Rel(base, value) + if err != nil { + return filepath.Base(value) + } + return relative +} + +func escapesRoot(value string) bool { + clean := filepath.Clean(value) + return clean == ".." || strings.HasPrefix(clean, ".."+string(os.PathSeparator)) +} + +func stringArg(args map[string]any, key string) (string, bool) { + value, ok := args[key].(string) + return value, ok +} + +func intArg(args map[string]any, key string, fallback int) int { + value, ok := args[key] + if !ok { + return fallback + } + switch number := value.(type) { + case int: + return number + case int64: + return int(number) + case float64: + return int(number) + default: + return fallback + } +} + +func errorResult(format string, values ...any) Result { + return Result{Output: "error: " + fmt.Sprintf(format, values...), IsError: true} +} + +func capOutput(output string, limit int) string { + characters := []rune(output) + if len(characters) <= limit { + return output + } + marker := []rune("\n…[truncated]") + if limit <= len(marker) { + return string(characters[:limit]) + } + return string(characters[:limit-len(marker)]) + string(marker) +} diff --git a/internal/mcp/builtin/fs_credentials_test.go b/internal/mcp/builtin/fs_credentials_test.go new file mode 100644 index 0000000..30facdf --- /dev/null +++ b/internal/mcp/builtin/fs_credentials_test.go @@ -0,0 +1,53 @@ +package builtin + +import ( + "context" + "os" + "path/filepath" + "strings" + "testing" +) + +func TestFilesystemToolsDenyRootlyCLICredentials(t *testing.T) { + home := t.TempDir() + t.Setenv("HOME", home) + credentialDir := filepath.Join(home, ".rootly-cli") + if err := os.MkdirAll(credentialDir, 0o700); err != nil { + t.Fatal(err) + } + for name, body := range map[string]string{ + "config.yaml": "api_key: rootly-test-secret", + "mcp-oauth.json": `{"access_token":"rootly-test-secret"}`, + } { + if err := os.WriteFile(filepath.Join(credentialDir, name), []byte(body), 0o600); err != nil { + t.Fatal(err) + } + } + + registry := testRegistry(t, home, nil) + readFile := mustLookup(t, registry, "local_read_file") + for _, path := range []string{".rootly-cli/config.yaml", ".rootly-cli/mcp-oauth.json"} { + result := readFile.Run(context.Background(), map[string]any{"path": path}) + if !result.IsError || !strings.Contains(result.Output, "path is denied") { + t.Errorf("read_file did not deny %q: %#v", path, result) + } + } + + grep := mustLookup(t, registry, "local_grep").Run(context.Background(), map[string]any{ + "pattern": "rootly-test-secret", + "output_mode": "content", + }) + if strings.Contains(grep.Output, "rootly-test-secret") || strings.Contains(grep.Output, ".rootly-cli") { + t.Errorf("grep exposed Rootly CLI credentials: %q", grep.Output) + } + + glob := mustLookup(t, registry, "local_glob").Run(context.Background(), map[string]any{"pattern": "**/*"}) + if strings.Contains(glob.Output, ".rootly-cli") || strings.Contains(glob.Output, "config.yaml") || strings.Contains(glob.Output, "mcp-oauth.json") { + t.Errorf("glob exposed Rootly CLI credential paths: %q", glob.Output) + } + + listing := mustLookup(t, registry, "local_list_dir").Run(context.Background(), map[string]any{"path": "."}) + if strings.Contains(listing.Output, ".rootly-cli") { + t.Errorf("list_dir exposed Rootly CLI credential directory: %q", listing.Output) + } +} diff --git a/internal/mcp/builtin/policy.go b/internal/mcp/builtin/policy.go new file mode 100644 index 0000000..b54e607 --- /dev/null +++ b/internal/mcp/builtin/policy.go @@ -0,0 +1,295 @@ +package builtin + +// Rules files use the prefix_rule syntax of openai/codex execpolicy +// (codex-rs/execpolicy, Apache-2.0) so policies can be shared between tools. + +import ( + "errors" + "fmt" + "os" + "path/filepath" + "slices" + "strings" + "sync" + + "go.starlark.net/starlark" + starsyntax "go.starlark.net/syntax" +) + +type Decision int + +const ( + DecisionNone Decision = iota + DecisionAllow + DecisionPrompt + DecisionForbidden +) + +const defaultRulesFile = "default.rules" + +// Rule matches argv whose leading tokens each equal one of the alternatives. +type Rule struct { + Pattern [][]string + Decision Decision + Justification string +} + +func (r Rule) Matches(argv []string) bool { + if len(r.Pattern) == 0 || len(argv) < len(r.Pattern) { + return false + } + for index, alternatives := range r.Pattern { + if !slices.Contains(alternatives, argv[index]) { + return false + } + } + return true +} + +// Policy holds user rules from ~/.rootly-cli/rules plus approvals granted +// for this session. The strictest matching decision wins. +type Policy struct { + mu sync.Mutex + dir string + rules []Rule + session []Rule + loadErr error +} + +func DefaultRulesDir() string { + home, err := os.UserHomeDir() + if err != nil || home == "" { + return "" + } + return filepath.Join(home, ".rootly-cli", "rules") +} + +// LoadPolicy reads every *.rules file in dir. A broken file is skipped and +// reported through LoadError; the built-in defaults still apply. +func LoadPolicy(dir string) *Policy { + policy := &Policy{dir: dir} + if dir == "" { + return policy + } + paths, _ := filepath.Glob(filepath.Join(dir, "*.rules")) + slices.Sort(paths) + var errs []error + for _, path := range paths { + rules, err := parseRulesFile(path) + if err != nil { + errs = append(errs, err) + continue + } + policy.rules = append(policy.rules, rules...) + } + policy.loadErr = errors.Join(errs...) + return policy +} + +func (p *Policy) Dir() string { return p.dir } + +func (p *Policy) LoadError() error { + if p == nil { + return nil + } + p.mu.Lock() + defer p.mu.Unlock() + return p.loadErr +} + +func (p *Policy) RuleCount() int { + if p == nil { + return 0 + } + p.mu.Lock() + defer p.mu.Unlock() + return len(p.rules) +} + +// Evaluate returns the strictest decision of the rules matching argv. +func (p *Policy) Evaluate(argv []string) (decision Decision, justification string) { + if p == nil { + return DecisionNone, "" + } + p.mu.Lock() + defer p.mu.Unlock() + for _, rule := range slices.Concat(p.rules, p.session) { + if rule.Decision > decision && rule.Matches(argv) { + decision, justification = rule.Decision, rule.Justification + } + } + return decision, justification +} + +func (p *Policy) AllowForSession(prefix []string) { + p.mu.Lock() + defer p.mu.Unlock() + p.session = append(p.session, allowRule(prefix)) +} + +// SaveAllow appends an allow rule for prefix to default.rules. +func (p *Policy) SaveAllow(prefix []string) error { + if p.dir == "" { + return errors.New("no home directory for rules") + } + if err := os.MkdirAll(p.dir, 0o700); err != nil { + return err + } + quoted := make([]string, len(prefix)) + for index, token := range prefix { + quoted[index] = starlark.String(token).String() + } + line := fmt.Sprintf("prefix_rule(pattern=[%s], decision=\"allow\")\n", strings.Join(quoted, ", ")) + path := filepath.Join(p.dir, defaultRulesFile) + if existing, err := os.ReadFile(path); err == nil && len(existing) > 0 && existing[len(existing)-1] != '\n' { + line = "\n" + line + } + file, err := os.OpenFile(path, os.O_APPEND|os.O_CREATE|os.O_WRONLY, 0o600) + if err != nil { + return err + } + if _, err := file.WriteString(line); err != nil { + _ = file.Close() + return err + } + if err := file.Close(); err != nil { + return err + } + p.mu.Lock() + defer p.mu.Unlock() + p.rules = append(p.rules, allowRule(prefix)) + return nil +} + +func allowRule(prefix []string) Rule { + pattern := make([][]string, len(prefix)) + for index, token := range prefix { + pattern[index] = []string{token} + } + return Rule{Pattern: pattern, Decision: DecisionAllow} +} + +func parseRulesFile(path string) ([]Rule, error) { + source, err := os.ReadFile(path) + if err != nil { + return nil, err + } + var rules []Rule + builtin := starlark.NewBuiltin("prefix_rule", func(_ *starlark.Thread, fn *starlark.Builtin, args starlark.Tuple, kwargs []starlark.Tuple) (starlark.Value, error) { + rule, err := unpackPrefixRule(fn, args, kwargs) + if err != nil { + return nil, err + } + rules = append(rules, rule) + return starlark.None, nil + }) + name := filepath.Base(path) + thread := &starlark.Thread{Name: name} + thread.SetMaxExecutionSteps(1_000_000) + if _, err := starlark.ExecFileOptions(&starsyntax.FileOptions{}, thread, name, source, starlark.StringDict{"prefix_rule": builtin}); err != nil { + var evalErr *starlark.EvalError + if errors.As(err, &evalErr) { + return nil, fmt.Errorf("%s: %w", name, err) + } + return nil, err + } + return rules, nil +} + +func unpackPrefixRule(fn *starlark.Builtin, args starlark.Tuple, kwargs []starlark.Tuple) (Rule, error) { + var ( + pattern *starlark.List + decision = "allow" + justification string + match *starlark.List + notMatch *starlark.List + ) + if err := starlark.UnpackArgs(fn.Name(), args, kwargs, + "pattern", &pattern, "decision?", &decision, "justification?", &justification, + "match?", &match, "not_match?", ¬Match); err != nil { + return Rule{}, err + } + rule := Rule{Justification: justification} + switch decision { + case "allow": + rule.Decision = DecisionAllow + case "prompt": + rule.Decision = DecisionPrompt + case "forbidden": + rule.Decision = DecisionForbidden + default: + return Rule{}, fmt.Errorf("prefix_rule: decision must be allow, prompt or forbidden, got %q", decision) + } + if pattern.Len() == 0 { + return Rule{}, errors.New("prefix_rule: pattern must not be empty") + } + for index := range pattern.Len() { + alternatives, err := tokenAlternatives(pattern.Index(index)) + if err != nil { + return Rule{}, fmt.Errorf("prefix_rule: pattern[%d]: %w", index, err) + } + rule.Pattern = append(rule.Pattern, alternatives) + } + if err := checkExamples(rule, match, true); err != nil { + return Rule{}, err + } + if err := checkExamples(rule, notMatch, false); err != nil { + return Rule{}, err + } + return rule, nil +} + +func tokenAlternatives(value starlark.Value) ([]string, error) { + switch value := value.(type) { + case starlark.String: + return []string{string(value)}, nil + case *starlark.List: + alternatives, err := stringList(value) + if err == nil && len(alternatives) == 0 { + err = errors.New("alternatives must not be empty") + } + return alternatives, err + } + return nil, fmt.Errorf("want a string or a list of strings, got %s", value.Type()) +} + +func stringList(list *starlark.List) ([]string, error) { + values := make([]string, 0, list.Len()) + for index := range list.Len() { + text, ok := starlark.AsString(list.Index(index)) + if !ok { + return nil, fmt.Errorf("want strings, got %s", list.Index(index).Type()) + } + values = append(values, text) + } + return values, nil +} + +func checkExamples(rule Rule, examples *starlark.List, want bool) error { + if examples == nil { + return nil + } + label := map[bool]string{true: "match", false: "not_match"}[want] + for index := range examples.Len() { + var argv []string + switch example := examples.Index(index).(type) { + case starlark.String: + commands, ok := shellCommands(string(example)) + if !ok || len(commands) != 1 { + return fmt.Errorf("prefix_rule: %s[%d]: not a plain command", label, index) + } + argv = commands[0] + case *starlark.List: + var err error + if argv, err = stringList(example); err != nil { + return fmt.Errorf("prefix_rule: %s[%d]: %w", label, index, err) + } + default: + return fmt.Errorf("prefix_rule: %s[%d]: want a string or a list of strings", label, index) + } + if rule.Matches(argv) != want { + return fmt.Errorf("prefix_rule: %s example %q does not hold", label, strings.Join(argv, " ")) + } + } + return nil +} diff --git a/internal/mcp/builtin/policy_test.go b/internal/mcp/builtin/policy_test.go new file mode 100644 index 0000000..5146996 --- /dev/null +++ b/internal/mcp/builtin/policy_test.go @@ -0,0 +1,613 @@ +package builtin + +import ( + "context" + "os" + "os/exec" + "path/filepath" + "runtime" + "slices" + "strings" + "testing" +) + +func TestBashKnownSafeCommandsRunWithoutApproval(t *testing.T) { + home := t.TempDir() + t.Setenv("HOME", home) + root := filepath.Join(home, "project") + if err := os.MkdirAll(filepath.Join(root, "src"), 0o755); err != nil { + t.Fatal(err) + } + for _, name := range []string{"README.md", ".env", filepath.Join("src", "main.go")} { + if err := os.WriteFile(filepath.Join(root, name), []byte("x"), 0o600); err != nil { + t.Fatal(err) + } + } + initTestRepo(t, root, "README.md", filepath.Join("src", "main.go")) + if err := os.Symlink(filepath.Join(home, "secret.txt"), filepath.Join(root, "escape")); err != nil { + t.Skipf("symlinks are unavailable: %v", err) + } + bash := mustLookup(t, testRegistryWithPolicy(t, root, LoadPolicy(filepath.Join(home, "rules"))), "local_bash") + safe := []string{ + "ls -la", "cat README.md", "head -n 5 src/main.go", "wc -l README.md src/main.go", + "git status", "git --no-pager log --oneline -5", "git diff HEAD~1 -- src", "git branch -a", + "git log --oneline | head -3", "grep -n TODO src/main.go 2>/dev/null", "sed -n 1,5p README.md", + "find . -name '*.go'", "pwd && whoami; uname -a", `echo "hello world"`, + } + for _, command := range safe { + if bash.NeedsApproval(map[string]any{"command": command}) { + t.Errorf("%q should run without approval", command) + } + } + asks := []string{ + "rm -rf src", "cat .env", "cat ../secret.txt", "cat /etc/passwd", "cat escape", "cat ~/.ssh/id_rsa", + "cat $HOME/.aws/credentials", "echo $(whoami)", "echo `id`", "cat *", "ls > out.txt", "cat README.md &", + "grep -r KEY .", "grep -rn KEY", "git push", "git -c core.pager=sh log", "git -C .. status", + "git diff --no-index /etc/passwd README.md", "git log --output=/tmp/x", "git branch -D main", + "find . -delete", "find . -exec rm {} ;", "sed -i s/a/b/ README.md", "base64 -o out README.md", + "FOO=1 ls", "(ls)", "for f in a; do ls; done", "ls | sh", "ls && curl example.com", "cat <err.txt", "cat < README.md", `cat '\etc\passwd'`, + "git show HEAD:.env", "git show HEAD:README.md", "git show :README.md", "git cat-file -p HEAD:README.md", + "git log -p -- .env", "git diff -- .env", "git diff .env", "git log -L1,2:README.md", "git log --full-diff -p -- src", + "git log --show-signature", "git log src -n 5", + } + for _, command := range asks { + if !bash.NeedsApproval(map[string]any{"command": command}) { + t.Errorf("%q must ask for approval", command) + } + } +} + +func TestRulesFileDecisionsApplyToBashAndCLIs(t *testing.T) { + home := t.TempDir() + t.Setenv("HOME", home) + root := filepath.Join(home, "project") + if err := os.MkdirAll(root, 0o755); err != nil { + t.Fatal(err) + } + dir := filepath.Join(home, "rules") + writeRules(t, dir, "default.rules", ` +prefix_rule(pattern=["kubectl", "rollout", ["status", "history"]], decision="allow", + match=[["kubectl", "rollout", "status", "deploy/api"], "kubectl rollout history"], + not_match=["kubectl rollout undo"]) +prefix_rule(pattern=["git", "push"], decision="forbidden", justification="Push from your own terminal.") +prefix_rule(pattern=["git", "status"], decision="prompt") +prefix_rule(pattern=["make", "test"]) +`) + policy := LoadPolicy(dir) + if err := policy.LoadError(); err != nil || policy.RuleCount() != 4 { + t.Fatalf("load: %v, %d rules", err, policy.RuleCount()) + } + lookPath := func(string) (string, error) { return "/bin/true", nil } + registry := testRegistryWithPolicy(t, root, policy) + kubectl := newCLITool(root, cliSpec{bin: "kubectl"}, lookPath, policy) + if kubectl.NeedsApproval(map[string]any{"args": []any{"rollout", "status", "deploy/api"}}) { + t.Error("allow rule should skip approval for kubectl rollout status") + } + if approval := kubectl.Approval(map[string]any{"args": []any{"rollout", "undo", "deploy/api"}}); !approval.Needs || + !slices.Equal(approval.Prefix, []string{"kubectl", "rollout", "undo"}) { + t.Errorf("unmatched mutating call: %+v", approval) + } + bash := mustLookup(t, registry, "local_bash") + if !bash.NeedsApproval(map[string]any{"command": "git status"}) { + t.Error("prompt rule must override the built-in safe list") + } + if approval := bash.Approval(map[string]any{"command": "git status"}); approval.Prefix != nil { + t.Errorf("a prompt rule must not be rememberable, got %v", approval.Prefix) + } + if bash.NeedsApproval(map[string]any{"command": "make test && git log -1"}) { + t.Error("allow rule plus safe command should run without approval") + } + pushed := bash.Approval(map[string]any{"command": "git log -1 && git push origin main"}) + if pushed.Forbidden != "Push from your own terminal." || pushed.Needs { + t.Fatalf("forbidden rule not applied: %+v", pushed) + } + result := bash.Run(context.Background(), map[string]any{"command": "git push origin main"}) + if !result.IsError || !strings.Contains(result.Output, "Push from your own terminal.") { + t.Fatalf("forbidden call ran: %+v", result) + } +} + +func TestRulesFileErrorsAreReportedAndSkipped(t *testing.T) { + dir := t.TempDir() + writeRules(t, dir, "a.rules", `prefix_rule(pattern=["ls"])`) + writeRules(t, dir, "b.rules", `prefix_rule(pattern=["git", "push"], match=["git pull"])`) + writeRules(t, dir, "c.rules", `prefix_rule(pattern=["rm"], decision="yolo")`) + writeRules(t, dir, "d.rules", `open("/etc/passwd")`) + policy := LoadPolicy(dir) + err := policy.LoadError() + for _, want := range []string{"b.rules", "does not hold", "c.rules", "decision must be", "d.rules", "undefined: open"} { + if err == nil || !strings.Contains(err.Error(), want) { + t.Errorf("load error %v should mention %q", err, want) + } + } + if policy.RuleCount() != 1 { + t.Fatalf("valid files should still load, got %d rules", policy.RuleCount()) + } +} + +func TestSessionAndSavedApprovals(t *testing.T) { + dir := filepath.Join(t.TempDir(), "rules") + policy := LoadPolicy(dir) + argv := []string{"terraform", "plan", "-out=x"} + policy.AllowForSession([]string{"terraform", "plan"}) + if decision, _ := policy.Evaluate(argv); decision != DecisionAllow { + t.Fatalf("session allow not applied: %v", decision) + } + if decision, _ := LoadPolicy(dir).Evaluate(argv); decision != DecisionNone { + t.Fatal("session approvals must not persist") + } + if err := policy.SaveAllow([]string{"gh", `pr"x`, "merge"}); err != nil { + t.Fatal(err) + } + reloaded := LoadPolicy(dir) + if err := reloaded.LoadError(); err != nil { + t.Fatal(err) + } + if decision, _ := reloaded.Evaluate([]string{"gh", `pr"x`, "merge", "12"}); decision != DecisionAllow { + t.Fatalf("saved rule not reloaded: %v", decision) + } + info, err := os.Stat(filepath.Join(dir, "default.rules")) + if err != nil { + t.Fatal(err) + } + if runtime.GOOS != "windows" && info.Mode().Perm() != 0o600 { + t.Fatalf("rules file mode: %v", info.Mode()) + } +} + +func TestSaveAllowAfterFileWithoutTrailingNewline(t *testing.T) { + dir := t.TempDir() + writeRules(t, dir, "default.rules", `prefix_rule(pattern=["make", "test"], decision="allow")`) + if err := LoadPolicy(dir).SaveAllow([]string{"kubectl", "rollout"}); err != nil { + t.Fatal(err) + } + reloaded := LoadPolicy(dir) + if err := reloaded.LoadError(); err != nil || reloaded.RuleCount() != 2 { + t.Fatalf("appended rule broke the file: %v, %d rules", err, reloaded.RuleCount()) + } +} + +// hermeticGit points git at empty user and system config so runner-wide +// settings (for example a git-lfs filter) can't change what the tests see. +func hermeticGit(t *testing.T) { + t.Helper() + if _, err := exec.LookPath("git"); err != nil { + t.Skip("git is not installed") + } + empty := t.TempDir() + global := filepath.Join(empty, "gitconfig") + if err := os.WriteFile(global, nil, 0o600); err != nil { + t.Fatal(err) + } + t.Setenv("XDG_CONFIG_HOME", empty) + t.Setenv("GIT_CONFIG_NOSYSTEM", "1") + t.Setenv("GIT_CONFIG_GLOBAL", global) +} + +func runGit(t *testing.T, root string, args ...string) string { + t.Helper() + cmd := exec.CommandContext(context.Background(), "git", append([]string{"-c", "user.email=e2e@example.com", "-c", "user.name=e2e"}, args...)...) + cmd.Dir = root + out, err := cmd.CombinedOutput() + if err != nil { + t.Fatalf("git %v: %v\n%s", args, err, out) + } + return strings.TrimSpace(string(out)) +} + +// initTestRepo commits every file under root twice, so HEAD~1 exists. +func initTestRepo(t *testing.T, root string, files ...string) { + t.Helper() + hermeticGit(t) + runGit(t, root, "init", "-q", "-b", "main") + runGit(t, root, append([]string{"add", "--"}, files...)...) + runGit(t, root, "commit", "-qm", "one") + if err := os.WriteFile(filepath.Join(root, files[0]), []byte("changed\n"), 0o600); err != nil { + t.Fatal(err) + } + runGit(t, root, "commit", "-qam", "two") +} + +func TestHardenSafeGitRewritesAutoApprovedReads(t *testing.T) { + root := t.TempDir() + if err := os.MkdirAll(filepath.Join(root, "src"), 0o755); err != nil { + t.Fatal(err) + } + if err := os.WriteFile(filepath.Join(root, "src", "a.go"), []byte("x"), 0o600); err != nil { + t.Fatal(err) + } + initTestRepo(t, root, filepath.Join("src", "a.go")) + config := "git " + gitHardeningConfig + cases := map[string]string{ + "git status": config + " status", + "git diff": config + " diff " + gitDiffHardening + " -- " + gitDeniedPathspecs, + "git diff HEAD~1 src": config + " diff " + gitDiffHardening + " HEAD~1 -- src " + gitDeniedPathspecs, + "git diff HEAD~1 -- src": config + " diff " + gitDiffHardening + " HEAD~1 -- src " + gitDeniedPathspecs, + "git --no-pager show --stat HEAD": config + " --no-pager show " + gitDiffHardening + " --stat HEAD -- " + gitDeniedPathspecs, + "git log --oneline | head -3": config + " log " + gitDiffHardening + " --oneline -- " + gitDeniedPathspecs + " | head -3", + "git log --follow -p -- src/a.go": config + " log " + gitDiffHardening + " --follow -p -- src/a.go", + "git branch -a && git log -1": config + " branch -a && " + config + " log " + gitDiffHardening + " -1 -- " + gitDeniedPathspecs, + "git push": "git push", + "rm -rf src; git diff": "rm -rf src; " + config + " diff " + gitDiffHardening + " -- " + gitDeniedPathspecs, + "echo $(git diff)": "echo $(" + config + " diff " + gitDiffHardening + " -- " + gitDeniedPathspecs + ")", + } + for script, want := range cases { + if got := hardenSafeGit(script, root); got != want { + t.Errorf("hardenSafeGit(%q)\n got %q\nwant %q", script, got, want) + } + } +} + +func TestGitDecidesWhichArgumentsAreRevisionsAndWhichArePaths(t *testing.T) { + home := t.TempDir() + root := filepath.Join(home, "repo") + if err := os.MkdirAll(filepath.Join(root, "src"), 0o755); err != nil { + t.Fatal(err) + } + for _, name := range []string{"README.md", "gone.go", filepath.Join("src", "main.go")} { + if err := os.WriteFile(filepath.Join(root, name), []byte("x\n"), 0o600); err != nil { + t.Fatal(err) + } + } + initTestRepo(t, root, "README.md", "gone.go", filepath.Join("src", "main.go")) + runGit(t, root, "branch", "src") + if err := os.Remove(filepath.Join(root, "gone.go")); err != nil { + t.Fatal(err) + } + readmeBlob := runGit(t, root, "rev-parse", "HEAD:README.md") + bash := mustLookup(t, testRegistryWithPolicy(t, root, LoadPolicy(filepath.Join(home, "rules"))), "local_bash") + config := "git " + gitHardeningConfig + rewrites := map[string]string{ + "git diff gone.go": config + " diff " + gitDiffHardening + " -- gone.go " + gitDeniedPathspecs, + "git log HEAD~1..HEAD": config + " log " + gitDiffHardening + " HEAD~1..HEAD -- " + gitDeniedPathspecs, + "git diff HEAD~1...HEAD": config + " diff " + gitDiffHardening + " HEAD~1...HEAD -- " + gitDeniedPathspecs, + "git diff src --": config + " diff " + gitDiffHardening + " src -- " + gitDeniedPathspecs, + "git diff -- src": config + " diff " + gitDiffHardening + " -- src " + gitDeniedPathspecs, + "git log -n 3 src/main.go": config + " log " + gitDiffHardening + " -n 3 -- src/main.go " + gitDeniedPathspecs, + "git log --oneline main..HEAD": config + " log " + gitDiffHardening + " --oneline main..HEAD -- " + gitDeniedPathspecs, + } + for command, want := range rewrites { + if bash.NeedsApproval(map[string]any{"command": command}) { + t.Errorf("%q should run without approval", command) + continue + } + if got := hardenSafeGit(command, root); got != want { + t.Errorf("hardenSafeGit(%q)\n got %q\nwant %q", command, got, want) + } + } + for _, command := range []string{"git diff gone.go", "git log HEAD~1..HEAD", "git diff src --"} { + if result := bash.Run(context.Background(), map[string]any{"command": command}); result.IsError { + t.Errorf("%q failed after hardening: %s", command, result.Output) + } + } + for _, command := range []string{ + "git diff src", "git log src", + "git show " + readmeBlob, "git show " + readmeBlob[:10], "git diff " + readmeBlob + ".." + readmeBlob, + "git show HEAD^{tree}", + } { + if !bash.NeedsApproval(map[string]any{"command": command}) { + t.Errorf("%q is ambiguous or names a blob/tree and must ask", command) + } + } +} + +func TestAutoApprovedGitReadsCannotRunRepoProgramsOrLeakDeniedFiles(t *testing.T) { + if _, err := exec.LookPath("git"); err != nil { + t.Skip("git is not installed") + } + if runtime.GOOS == "windows" { + t.Skip("drivers below use POSIX shell commands") + } + home := t.TempDir() + t.Setenv("HOME", home) + root := filepath.Join(home, "repo") + if err := os.MkdirAll(root, 0o755); err != nil { + t.Fatal(err) + } + git := func(args ...string) { + t.Helper() + cmd := exec.CommandContext(context.Background(), "git", append([]string{"-c", "user.email=e2e@example.com", "-c", "user.name=e2e"}, args...)...) + cmd.Dir = root + cmd.Env = append(os.Environ(), "GIT_CONFIG_NOSYSTEM=1", "GIT_CONFIG_GLOBAL="+os.DevNull) + if out, err := cmd.CombinedOutput(); err != nil { + t.Fatalf("git %v: %v\n%s", args, err, out) + } + } + write := func(name, body string) { + t.Helper() + if err := os.WriteFile(filepath.Join(root, name), []byte(body), 0o600); err != nil { + t.Fatal(err) + } + } + marker := filepath.Join(home, "pwned") + git("init", "-q") + write("README.md", "one\n") + write(".env", "TOKEN=first\n") + write(".gitattributes", "*.md diff=evil\n") + git("add", "README.md", ".env", ".gitattributes") + git("commit", "-qm", "init") + write("README.md", "one\ntwo\n") + write(".env", "TOKEN=leaked-secret\n") + git("commit", "-qam", "two") + write("README.md", "one\ntwo\nthree\n") + git("config", "diff.evil.textconv", "touch "+marker+"; cat") + git("config", "diff.external", "touch "+marker+";") + git("config", "core.fsmonitor", "touch "+marker+"; true") + bash := mustLookup(t, testRegistryWithPolicy(t, root, LoadPolicy(filepath.Join(home, "rules"))), "local_bash") + for _, command := range []string{"git diff", "git show HEAD", "git log -p -2", "git status"} { + args := map[string]any{"command": command} + if bash.NeedsApproval(args) { + t.Fatalf("%q should run without approval", command) + } + result := bash.Run(context.Background(), args) + if result.IsError || strings.Contains(result.Output, "leaked-secret") || strings.Contains(result.Output, "TOKEN=") { + t.Fatalf("%q: %+v", command, result) + } + if _, err := os.Stat(marker); err == nil { + t.Fatalf("%q ran a repo-configured program", command) + } + } + gitOut := func(args ...string) string { + t.Helper() + cmd := exec.CommandContext(context.Background(), "git", args...) + cmd.Dir = root + out, err := cmd.Output() + if err != nil { + t.Fatalf("git %v: %v", args, err) + } + return strings.TrimSpace(string(out)) + } + envBlob, oldEnvBlob := gitOut("rev-parse", "HEAD:.env"), gitOut("rev-parse", "HEAD~1:.env") + git("tag", "envblob", envBlob) + for _, command := range []string{ + "git show " + envBlob, "git show " + envBlob[:12], "git diff " + oldEnvBlob + " " + envBlob, + "git diff " + oldEnvBlob + ".." + envBlob, "git show HEAD^{tree}", "git show envblob", "git log -p " + envBlob, + "git show v9.9.9", + } { + if !bash.NeedsApproval(map[string]any{"command": command}) { + t.Errorf("%q names a blob or tree and must ask", command) + } + } + for _, command := range []string{"git show HEAD~1", "git diff HEAD~1..HEAD", "git log -n 2 -p"} { + if bash.NeedsApproval(map[string]any{"command": command}) { + t.Errorf("%q names commits only and should not ask", command) + } + } + git("config", "filter.evil.clean", "touch "+marker+"; cat") + write(".gitattributes", "*.md diff=evil filter=evil\n") + for _, command := range []string{"git diff", "git status"} { + if !bash.NeedsApproval(map[string]any{"command": command}) { + t.Errorf("%q must ask when the repo config defines a clean filter", command) + } + } + if bash.NeedsApproval(map[string]any{"command": "git log -1"}) { + t.Error("git log does not run filters and should not ask") + } +} + +func TestRememberPrefix(t *testing.T) { + cases := map[string][]string{ + "kubectl rollout undo deploy/api": {"kubectl", "rollout", "undo"}, + "gh pr merge 12": {"gh", "pr", "merge"}, + "terraform apply": {"terraform", "apply"}, + "npm test": {"npm", "test"}, + "psql -c select": nil, + "rm -rf /": nil, + "ls": nil, + "./deploy.sh prod": nil, + "curl example.com": nil, + "sudo rm -rf x": nil, + "bash deploy": nil, + "python manage": nil, + "env kubectl delete": nil, + "docker run alpine": nil, + } + for command, want := range cases { + if got := rememberPrefix(strings.Fields(command)); !slices.Equal(got, want) { + t.Errorf("%q: got %v, want %v", command, got, want) + } + } +} + +func testRegistryWithPolicy(t *testing.T, root string, policy *Policy) *Registry { + t.Helper() + registry, err := NewRegistryWithPolicy(root, func(string) (string, error) { return "", os.ErrNotExist }, policy) + if err != nil { + t.Fatal(err) + } + return registry +} + +func writeRules(t *testing.T, dir, name, source string) { + t.Helper() + if err := os.MkdirAll(dir, 0o700); err != nil { + t.Fatal(err) + } + if err := os.WriteFile(filepath.Join(dir, name), []byte(source), 0o600); err != nil { + t.Fatal(err) + } +} + +func TestSymlinkedWorkspaceRootIsComparedResolved(t *testing.T) { + base := t.TempDir() + target := filepath.Join(base, "target") + if err := os.MkdirAll(filepath.Join(target, "src"), 0o755); err != nil { + t.Fatal(err) + } + for name, body := range map[string]string{ + "README.md": "x\n", ".env": "TOKEN=x\n", filepath.Join("src", "a.go"): "x\n", filepath.Join("..", "outside.txt"): "x\n", + } { + if err := os.WriteFile(filepath.Join(target, name), []byte(body), 0o600); err != nil { + t.Fatal(err) + } + } + alias := filepath.Join(base, "alias") + if err := os.Symlink(target, alias); err != nil { + t.Skipf("symlinks are unavailable: %v", err) + } + if err := os.Symlink(filepath.Join(base, "outside.txt"), filepath.Join(target, "escape")); err != nil { + t.Skipf("symlinks are unavailable: %v", err) + } + initTestRepo(t, target, "README.md", filepath.Join("src", "a.go")) + for _, root := range []string{alias, target} { + inside := []string{"README.md", filepath.Join("src", "a.go"), "missing.txt", filepath.Join(alias, "README.md"), filepath.Join(target, "README.md")} + for _, candidate := range inside { + if !safeReadPath(root, candidate) { + t.Errorf("root %s: %q is inside the workspace", root, candidate) + } + } + outside := []string{ + "escape", ".env", filepath.Join("..", "outside.txt"), filepath.Join(base, "outside.txt"), + filepath.Join(alias, ".env"), filepath.Join(alias, "escape"), filepath.Join(alias, "..", "outside.txt"), + } + for _, candidate := range outside { + if safeReadPath(root, candidate) { + t.Errorf("root %s: %q must stay outside the workspace", root, candidate) + } + } + config := "git " + gitHardeningConfig + want := config + " log " + gitDiffHardening + " --follow -p -- src/a.go" + if got := hardenSafeGit("git log --follow -p -- src/a.go", root); got != want { + t.Errorf("root %s: hardenSafeGit\n got %q\nwant %q", root, got, want) + } + if !isKnownSafeCommand([]string{"cat", "README.md"}, root) || isKnownSafeCommand([]string{"cat", "escape"}, root) { + t.Errorf("root %s: cat classification changed through the symlinked root", root) + } + } +} + +func TestWorktreeConfigFiltersAsk(t *testing.T) { + if runtime.GOOS == "windows" { + t.Skip("filter below uses a POSIX shell command") + } + root := t.TempDir() + if err := os.WriteFile(filepath.Join(root, "README.md"), []byte("x\n"), 0o600); err != nil { + t.Fatal(err) + } + initTestRepo(t, root, "README.md") + runGit(t, root, "config", "extensions.worktreeConfig", "true") + if repoConfigRunsFilters([]string{"diff"}, root) { + t.Fatal("no filter configured yet") + } + runGit(t, root, "config", "--worktree", "filter.evil.clean", "touch pwned; cat") + for _, command := range [][]string{{"git", "diff"}, {"git", "status"}} { + if isKnownSafeCommand(command, root) { + t.Errorf("%v must ask when config.worktree defines a clean filter", command) + } + } +} + +func TestFileSafeCommandsCannotWriteOrReadOutsideNames(t *testing.T) { + root := t.TempDir() + for _, name := range []string{"in.txt", "out.txt", "list"} { + if err := os.WriteFile(filepath.Join(root, name), []byte("x\n"), 0o600); err != nil { + t.Fatal(err) + } + } + for _, argv := range [][]string{ + {"uniq", "in.txt"}, {"uniq", "-c", "in.txt"}, {"uniq", "-f", "1", "-s", "2", "in.txt"}, {"uniq"}, {"wc", "-l", "in.txt"}, + } { + if !isKnownSafeCommand(argv, root) { + t.Errorf("%v only reads and should run unasked", argv) + } + } + for _, argv := range [][]string{ + {"uniq", "in.txt", "out.txt"}, {"uniq", "-", "out.txt"}, {"uniq", "-c", "in.txt", "missing.txt"}, {"uniq", "--", "in.txt", "out.txt"}, + {"wc", "--files0-from=list"}, {"wc", "--files0-from", "list"}, {"find", ".", "-files0-from", "list"}, + } { + if isKnownSafeCommand(argv, root) { + t.Errorf("%v writes or reads names it doesn't show and must ask", argv) + } + } +} + +func TestGitBlobSpecsAskEvenForDeniedPaths(t *testing.T) { + root := t.TempDir() + if err := os.MkdirAll(filepath.Join(root, "sub"), 0o755); err != nil { + t.Fatal(err) + } + for _, name := range []string{"README.md", ".env", filepath.Join("sub", ".env")} { + if err := os.WriteFile(filepath.Join(root, name), []byte("TOKEN=x\n"), 0o600); err != nil { + t.Fatal(err) + } + } + initTestRepo(t, root, "README.md", ".env", filepath.Join("sub", ".env")) + for _, spec := range []string{"HEAD:.env", ":.env", ":0:.env", "HEAD:sub/.env", "HEAD:./.env", "HEAD~1:.env", "HEAD:README.md"} { + for _, sub := range []string{"show", "log", "diff"} { + if isKnownSafeCommand([]string{"git", sub, spec}, root) { + t.Errorf("git %s %s names a blob, which pathspec excludes can't filter, and must ask", sub, spec) + } + } + } +} + +func TestForbiddenRulesSurviveRedirectsAndComplexSyntax(t *testing.T) { + home := t.TempDir() + rules := filepath.Join(home, "rules") + if err := os.MkdirAll(rules, 0o700); err != nil { + t.Fatal(err) + } + body := "prefix_rule(pattern=[\"git\", \"push\"], decision=\"forbidden\", justification=\"no pushes\")\n" + if err := os.WriteFile(filepath.Join(rules, "default.rules"), []byte(body), 0o600); err != nil { + t.Fatal(err) + } + marker := filepath.Join(home, "ran") + bash := mustLookup(t, testRegistryWithPolicy(t, home, LoadPolicy(rules)), "local_bash") + for _, command := range []string{ + "git push", "git push > push.log", "git push 2>&1 | tee push.log", "echo $(git push)", "x=1 git push", + "sudo git push", "env A=1 git push", "/usr/bin/git push", "bash -c 'git push'", "eval git push", + "true && { git push; }", "(git push)", "for i in 1; do git push; done", "git push &", `git "push"`, + "git push ((", + } { + approval := bash.Approval(map[string]any{"command": command}) + if approval.Forbidden != "no pushes" { + t.Errorf("%q lost its forbidden verdict: %+v", command, approval) + } + result := bash.Run(context.Background(), map[string]any{"command": command + "; touch " + marker}) + if !result.IsError || !strings.Contains(result.Output, "Blocked by a rule") { + t.Errorf("%q ran despite a forbidden rule: %s", command, result.Output) + } + } + if _, err := os.Stat(marker); err == nil { + t.Fatal("a forbidden script ran") + } + if approval := bash.Approval(map[string]any{"command": "git status > /tmp/x"}); approval.Forbidden != "" || !approval.Needs { + t.Fatalf("an unrelated redirect should ask, not be forbidden: %+v", approval) + } +} + +func TestGitRevisionSetsAreCommits(t *testing.T) { + root := t.TempDir() + if err := os.WriteFile(filepath.Join(root, "README.md"), []byte("x\n"), 0o600); err != nil { + t.Fatal(err) + } + initTestRepo(t, root, "README.md") + bash := mustLookup(t, testRegistryWithPolicy(t, root, LoadPolicy(filepath.Join(t.TempDir(), "rules"))), "local_bash") + config := "git " + gitHardeningConfig + for command, want := range map[string]string{ + "git log --oneline HEAD^!": config + " log " + gitDiffHardening + " --oneline HEAD^! -- " + gitDeniedPathspecs, + "git log -1 HEAD^@": config + " log " + gitDiffHardening + " -1 HEAD^@ -- " + gitDeniedPathspecs, + "git log HEAD^-": config + " log " + gitDiffHardening + " HEAD^- -- " + gitDeniedPathspecs, + "git diff HEAD^-1": config + " diff " + gitDiffHardening + " HEAD^-1 -- " + gitDeniedPathspecs, + } { + if bash.NeedsApproval(map[string]any{"command": command}) { + t.Errorf("%q names commits and should run unasked", command) + continue + } + if got := hardenSafeGit(command, root); got != want { + t.Errorf("hardenSafeGit(%q)\n got %q\nwant %q", command, got, want) + } + } + for _, command := range []string{"git log --oneline HEAD^!", "git log -1 HEAD^@"} { + result := bash.Run(context.Background(), map[string]any{"command": command}) + if result.IsError || strings.TrimSpace(result.Output) == "" { + t.Errorf("%q returned no commits after hardening: %s", command, result.Output) + } + } + for _, command := range []string{"git log nosuchref", "git log ^!", "git log HEAD^!..HEAD", "git show nosuch^@"} { + if !bash.NeedsApproval(map[string]any{"command": command}) { + t.Errorf("%q is neither a commit nor a known path and must ask", command) + } + } +} diff --git a/internal/mcp/builtin/registry.go b/internal/mcp/builtin/registry.go new file mode 100644 index 0000000..a22663d --- /dev/null +++ b/internal/mcp/builtin/registry.go @@ -0,0 +1,137 @@ +package builtin + +import ( + "context" + "encoding/json" + "os/exec" + "path/filepath" + "strings" +) + +type Result struct { + Output string + IsError bool +} + +type Tool interface { + Name() string + Description() string + Schema() json.RawMessage + AdvertisedReadOnly() bool + NeedsApproval(map[string]any) bool + Approval(map[string]any) Approval + Run(context.Context, map[string]any) Result +} + +// Approval is the local decision for one call. Prefix is the command prefix +// the user may allow for the session or save as a rule; nil means the call +// can only be approved once. ReadOnly means the call only reads on its own +// merits, not because a rule or approval allows it, so it may run alongside +// other reads. +type Approval struct { + Needs bool + Forbidden string + Prefix []string + ReadOnly bool +} + +type LookPathFunc func(string) (string, error) + +type Registry struct { + tools []Tool + byName map[string]Tool + policy *Policy +} + +type localTool struct { + name string + description string + schema json.RawMessage + readOnly bool + needsApproval func(map[string]any) bool + approve func(map[string]any) Approval + run func(context.Context, map[string]any) Result +} + +func (t localTool) Name() string { return t.name } +func (t localTool) Description() string { return t.description } +func (t localTool) Schema() json.RawMessage { return append(json.RawMessage(nil), t.schema...) } +func (t localTool) AdvertisedReadOnly() bool { return t.readOnly } +func (t localTool) NeedsApproval(args map[string]any) bool { return t.Approval(args).Needs } + +func (t localTool) Approval(args map[string]any) Approval { + if t.approve != nil { + return t.approve(args) + } + needs := t.needsApproval(args) + return Approval{Needs: needs, ReadOnly: !needs} +} + +func (t localTool) Run(ctx context.Context, args map[string]any) Result { + if t.approve != nil { + if forbidden := t.approve(args).Forbidden; forbidden != "" { + return errorResult("Blocked by a rule in ~/.rootly-cli/rules: %s", forbidden) + } + } + return t.run(ctx, args) +} + +func NewRegistry(root string, lookPath LookPathFunc) (*Registry, error) { + return NewRegistryWithPolicy(root, lookPath, LoadPolicy(DefaultRulesDir())) +} + +func NewRegistryWithPolicy(root string, lookPath LookPathFunc, policy *Policy) (*Registry, error) { + if policy == nil { + policy = &Policy{} + } + if lookPath == nil { + lookPath = exec.LookPath + } + canonicalRoot, err := canonicalizeRoot(root) + if err != nil { + return nil, err + } + registry := &Registry{byName: map[string]Tool{}, policy: policy} + registry.add(newReadFileTool(canonicalRoot)) + registry.add(newGrepTool(canonicalRoot)) + registry.add(newGlobTool(canonicalRoot)) + registry.add(newListDirTool(canonicalRoot)) + for _, spec := range cliRegistry { + if _, err := lookPath(spec.bin); err == nil { + registry.add(newCLITool(canonicalRoot, spec, lookPath, policy)) + } + } + registry.add(newBashTool(canonicalRoot, policy)) + return registry, nil +} + +func (r *Registry) add(tool Tool) { + r.tools = append(r.tools, tool) + r.byName[tool.Name()] = tool +} + +func (r *Registry) Tools() []Tool { + return append([]Tool(nil), r.tools...) +} + +func (r *Registry) Policy() *Policy { return r.policy } + +func (r *Registry) Lookup(name string) (Tool, bool) { + tool, ok := r.byName[name] + return tool, ok +} + +func canonicalizeRoot(root string) (string, error) { + if strings.TrimSpace(root) == "" { + var err error + root, err = filepath.Abs(".") + if err != nil { + return "", err + } + } + absolute, err := filepath.Abs(root) + if err != nil { + return "", err + } + return filepath.EvalSymlinks(absolute) +} diff --git a/internal/mcp/builtin/safety.go b/internal/mcp/builtin/safety.go new file mode 100644 index 0000000..c397292 --- /dev/null +++ b/internal/mcp/builtin/safety.go @@ -0,0 +1,693 @@ +package builtin + +// The read-only command rules below follow google-gemini/gemini-cli +// (packages/core/src/sandbox/utils/commandSafety.ts, Apache-2.0, +// Copyright 2025 Google LLC), tightened for Rootly: no recursion into +// denied files, no git global options, and no shell expansion at all. + +import ( + "context" + "errors" + "os" + "os/exec" + "path/filepath" + "regexp" + "slices" + "strings" + "time" + "unicode" + + "mvdan.cc/sh/v3/syntax" +) + +// shellCommands splits a bash script into the argv of each simple command. +// It returns ok=false for anything it cannot read literally: expansions, +// globs, subshells, loops, assignments, background jobs or redirections +// other than 2>&1 and >/dev/null. +func shellCommands(script string) ([][]string, bool) { + file, err := syntax.NewParser(syntax.Variant(syntax.LangBash)).Parse(strings.NewReader(script), "") + if err != nil || len(file.Stmts) == 0 { + return nil, false + } + var commands [][]string + for _, stmt := range file.Stmts { + if !collectStmt(stmt, &commands) { + return nil, false + } + } + return commands, len(commands) > 0 +} + +func collectStmt(stmt *syntax.Stmt, commands *[][]string) bool { + if stmt == nil || stmt.Negated || stmt.Background || stmt.Coprocess || stmt.Disown { + return false + } + for _, redirect := range stmt.Redirs { + if !harmlessRedirect(redirect) { + return false + } + } + switch command := stmt.Cmd.(type) { + case *syntax.CallExpr: + if len(command.Assigns) > 0 || len(command.Args) == 0 { + return false + } + argv := make([]string, 0, len(command.Args)) + for _, word := range command.Args { + value, ok := literalWord(word) + if !ok { + return false + } + argv = append(argv, value) + } + *commands = append(*commands, argv) + return true + case *syntax.BinaryCmd: + switch command.Op { + case syntax.AndStmt, syntax.OrStmt, syntax.Pipe, syntax.PipeAll: + return collectStmt(command.X, commands) && collectStmt(command.Y, commands) + } + } + return false +} + +func harmlessRedirect(redirect *syntax.Redirect) bool { + if redirect == nil || redirect.Hdoc != nil { + return false + } + fd := "" + if redirect.N != nil { + fd = redirect.N.Value + } + target, ok := literalWord(redirect.Word) + if !ok { + return false + } + switch redirect.Op { + case syntax.DplOut: + return fd == "2" && target == "1" + case syntax.RdrOut, syntax.AppOut: + return (fd == "" || fd == "1" || fd == "2") && target == "/dev/null" + case syntax.RdrAll, syntax.AppAll: + return fd == "" && target == "/dev/null" + } + return false +} + +func literalWord(word *syntax.Word) (string, bool) { + if word == nil || len(word.Parts) == 0 { + return "", false + } + var value strings.Builder + for index, part := range word.Parts { + switch part := part.(type) { + case *syntax.Lit: + if strings.ContainsAny(part.Value, "*?[]{}\\$`") || tildeExpands(part.Value, index == 0) { + return "", false + } + value.WriteString(part.Value) + case *syntax.SglQuoted: + if part.Dollar { + return "", false + } + value.WriteString(part.Value) + case *syntax.DblQuoted: + if part.Dollar { + return "", false + } + for _, inner := range part.Parts { + lit, ok := inner.(*syntax.Lit) + if !ok || strings.ContainsAny(lit.Value, "\\$`") { + return "", false + } + value.WriteString(lit.Value) + } + default: + return "", false + } + } + return value.String(), true +} + +// tildeExpands reports whether bash would expand ~ in this unquoted literal. +func tildeExpands(value string, wordStart bool) bool { + return (wordStart && strings.HasPrefix(value, "~")) || strings.Contains(value, "=~") || strings.Contains(value, ":~") +} + +var ( + plainSafeCommands = []string{"echo", "expr", "false", "id", "pwd", "seq", "true", "uname", "which", "whoami", "numfmt", "tr"} + fileSafeCommands = []string{"cat", "cut", "head", "ls", "nl", "paste", "rev", "stat", "tail", "tac", "uniq", "wc"} + sedPrintArg = regexp.MustCompile(`^\d+(,\d+)?p$`) +) + +// isKnownSafeCommand reports whether argv only reads files under root. +func isKnownSafeCommand(argv []string, root string) bool { + if len(argv) == 0 { + return false + } + root = resolveRoot(root) + name, args := argv[0], argv[1:] + switch { + case slices.Contains(plainSafeCommands, name): + return true + case slices.Contains(fileSafeCommands, name): + return fileCommandOnlyReads(name, args) && pathArgsStayInRoot(args, root) + } + switch name { + case "grep": + return !grepRecurses(args) && pathArgsStayInRoot(args, root) + case "base64": + return !hasFlag(args, "-o", "--output") && pathArgsStayInRoot(args, root) + case "find": + return findIsReadOnly(args) && pathArgsStayInRoot(args, root) + case "sed": + return len(args) >= 2 && len(args) <= 3 && args[0] == "-n" && sedPrintArg.MatchString(args[1]) && + pathArgsStayInRoot(args[2:], root) + case "git": + return gitIsReadOnly(args, root) && !repoConfigRunsFilters(args, root) + } + return false +} + +// fileCommandOnlyReads rejects the operands and flags of fileSafeCommands +// that write a file or read paths from a file: uniq's second operand is its +// output, and wc --files0-from reads names that pathArgsStayInRoot never sees. +func fileCommandOnlyReads(name string, args []string) bool { + switch name { + case "uniq": + return uniqOperands(args) <= 1 + case "wc": + return !hasFlag(args, "--files0-from") + } + return true +} + +func uniqOperands(args []string) int { + operands := 0 + for index := 0; index < len(args); index++ { + arg := args[index] + switch { + case arg == "--": + return operands + len(args) - index - 1 + case arg == "-f" || arg == "-s" || arg == "-w": + index++ + case arg == "-" || !strings.HasPrefix(arg, "-"): + operands++ + } + } + return operands +} + +func hasFlag(args []string, flags ...string) bool { + for _, arg := range args { + for _, flag := range flags { + if arg == flag || strings.HasPrefix(arg, flag+"=") || (len(flag) == 2 && !strings.HasPrefix(arg, "--") && strings.HasPrefix(arg, flag)) { + return true + } + } + } + return false +} + +func grepRecurses(args []string) bool { + for _, arg := range args { + switch { + case arg == "--": + return false + case strings.HasPrefix(arg, "--"): + if arg == "--recursive" || arg == "--dereference-recursive" || strings.HasPrefix(arg, "--directories") || + strings.HasPrefix(arg, "--include") || strings.HasPrefix(arg, "--exclude") { + return true + } + case strings.HasPrefix(arg, "-") && strings.ContainsAny(arg, "rRd"): + return true + } + } + return false +} + +func findIsReadOnly(args []string) bool { + unsafe := []string{"-exec", "-execdir", "-ok", "-okdir", "-delete", "-fls", "-fprint", "-fprint0", "-fprintf", "-files0-from"} + for _, arg := range args { + if slices.Contains(unsafe, arg) { + return false + } + } + return true +} + +// pathArgsStayInRoot treats every argument (and every --flag=value) as a +// possible path: none may leave root or name a denied file. +func pathArgsStayInRoot(args []string, root string) bool { + for _, arg := range args { + candidate := arg + if strings.HasPrefix(arg, "-") { + _, value, ok := strings.Cut(arg, "=") + if !ok { + continue + } + candidate = value + } + if candidate == "" { + continue + } + if !safeReadPath(root, candidate) { + return false + } + } + return true +} + +// safeReadPath reports whether candidate stays in the workspace and misses +// the denied list. An existing path is judged by where it really points, so +// a symlink can't escape and an alias of root (macOS /var -> /private/var) +// still counts as inside. +func safeReadPath(root, candidate string) bool { + root = resolveRoot(root) + path := candidate + switch { + case filepath.IsAbs(path): + case filepath.VolumeName(path) != "" || strings.HasPrefix(path, "/") || strings.HasPrefix(path, `\`): + // Rooted or drive-relative on Windows: never inside the workspace. + return false + default: + path = filepath.Join(root, path) + } + path = filepath.Clean(path) + if deniedPath(path) { + return false + } + if _, err := os.Lstat(path); err != nil { + return isWithin(root, path) + } + resolved, err := filepath.EvalSymlinks(path) + return err == nil && isWithin(root, resolved) && !deniedPath(resolved) +} + +// resolveRoot returns the workspace root as an absolute path with symlinks +// resolved, so it compares equal to resolved candidates. +func resolveRoot(root string) string { + absolute, err := filepath.Abs(root) + if err != nil { + return filepath.Clean(root) + } + if resolved, err := filepath.EvalSymlinks(absolute); err == nil { + return resolved + } + return absolute +} + +var ( + gitUnsafeFlags = []string{ + "--output", "--ext-diff", "--textconv", "--exec", "--paginate", "--no-index", "--full-diff", "--show-signature", + "--open-files-in-pager", + } + gitDiffCommands = []string{"diff", "log", "show"} +) + +// gitIsReadOnly allows status/log/diff/show/branch reads. Every revision or +// path must stay in the workspace and miss the denied list; blob specs +// (rev:path) and pathspec magic ask because they can name any file. +func gitIsReadOnly(args []string, root string) bool { + index := 0 + for index < len(args) && strings.HasPrefix(args[index], "-") { + if args[index] != "--no-pager" { + return false + } + index++ + } + if index >= len(args) { + return false + } + subcommand, rest := args[index], args[index+1:] + pathsOnly := false + for _, arg := range rest { + switch { + case pathsOnly || !strings.HasPrefix(arg, "-"): + if strings.Contains(arg, ":") || !safeReadPath(root, arg) { + return false + } + case arg == "--": + pathsOnly = true + case slices.Contains(gitUnsafeFlags, arg) || strings.HasPrefix(arg, "--output=") || strings.HasPrefix(arg, "--exec=") || + strings.HasPrefix(arg, "-L"): + return false + } + } + switch subcommand { + case "status": + return true + case "log", "diff", "show": + _, ok := gitPathStart(rest, root) + return ok + case "branch": + return gitBranchIsReadOnly(rest) + } + return false +} + +// gitPathStart finds where paths begin in the arguments after a diff/log/show +// subcommand: at "--", or at the first positional git doesn't resolve to a +// commit. ok is false (ask) when a positional names a blob or tree, is both +// a commit and an existing path with no "--", or can't be classified, and +// when options follow a path, which is ambiguous to rewrite. +func gitPathStart(rest []string, root string) (start int, ok bool) { + start = -1 + for index := 0; index < len(rest); index++ { + arg := rest[index] + switch { + case arg == "--": + return index, start < 0 + case start >= 0: + if strings.HasPrefix(arg, "-") { + return -1, false + } + case strings.HasPrefix(arg, "-"): + if slices.Contains(gitValueFlags, arg) { + index++ + } + default: + switch gitClassifyRev(root, arg) { + case gitRevCommit: + if _, err := os.Lstat(filepath.Join(root, arg)); err == nil && !slices.Contains(rest[index:], "--") { + return -1, false + } + case gitRevPath: + start = index + default: + return -1, false + } + } + } + return start, true +} + +// gitValueFlags take their value as the next argument, which is neither a +// revision nor a path. +var gitValueFlags = []string{ + "-n", "--max-count", "--skip", "--since", "--after", "--until", "--before", "--author", + "--committer", "--grep", "--format", "--pretty", "--date", "--diff-filter", "-G", "-S", +} + +type gitRevKind int + +const ( + gitRevPath gitRevKind = iota + gitRevCommit + gitRevContent + gitRevUnknown +) + +// gitClassifyRev asks git what a positional means: a commit (a range or a +// revision set), a blob or tree object, whose content pathspec excludes +// can't filter, a path git knows, or unknown, which asks rather than guess. +func gitClassifyRev(root, arg string) gitRevKind { + sides := []string{arg} + isRange := strings.Contains(arg, "..") + if isRange { + sides = strings.Split(strings.ReplaceAll(arg, "...", ".."), "..") + } + for _, side := range sides { + if side == "" { + continue + } + if base := gitRevSetSuffix.ReplaceAllString(side, ""); base != side { + if base == "" || isRange { + return gitRevUnknown + } + side = base + } + ok, err := gitSucceeds(root, "rev-parse", "--verify", "--quiet", "--end-of-options", side+"^{commit}") + switch { + case err != nil: + return gitRevUnknown + case ok: + continue + } + output, err := gitOutput(root, "cat-file", "-t", side+"^{}") + switch { + case err != nil: + return gitRevUnknown + case output == "blob" || output == "tree": + return gitRevContent + case isRange || !gitKnowsPath(root, side): + return gitRevUnknown + } + return gitRevPath + } + return gitRevCommit +} + +// gitRevSetSuffix matches the revision-set notations HEAD^!, HEAD^@ and +// HEAD^-N, which name commits but don't peel with ^{commit}. +var gitRevSetSuffix = regexp.MustCompile(`\^(!|@|-\d*)$`) + +// gitKnowsPath reports whether arg is clearly a path: it exists in the +// worktree or git tracks it (e.g. a deleted file not yet committed). +func gitKnowsPath(root, arg string) bool { + if _, err := os.Lstat(filepath.Join(root, arg)); err == nil { + return true + } + ok, err := gitSucceeds(root, "ls-files", "--error-unmatch", "--", ":(literal)"+arg) + return err == nil && ok +} + +// gitSucceeds runs a hardened git lookup in root. ok is false when git exits +// non-zero; err is set when git can't give an answer (killed, timed out). +func gitSucceeds(root string, args ...string) (ok bool, err error) { + _, ok, err = runGitLookup(root, args) + return ok, err +} + +func gitOutput(root string, args ...string) (string, error) { + output, _, err := runGitLookup(root, args) + return output, err +} + +func runGitLookup(root string, args []string) (output string, ok bool, err error) { + ctx, cancel := context.WithTimeout(context.Background(), 3*time.Second) + defer cancel() + cmd := exec.CommandContext(ctx, "git", append([]string{"-c", "core.fsmonitor=false"}, args...)...) + cmd.Dir = root + raw, runErr := cmd.Output() + var exitError *exec.ExitError + switch { + case runErr == nil: + return strings.TrimSpace(string(raw)), true, nil + case errors.Is(runErr, exec.ErrNotFound): + return "", false, nil + case errors.As(runErr, &exitError) && exitError.ExitCode() > 0: + return "", false, nil + } + return "", false, runErr +} + +// repoConfigRunsFilters reports whether repository or worktree config (not +// the user's own global or system config) defines clean/smudge filters, which git diff and +// status run on working-tree files and no flag can switch off. +func repoConfigRunsFilters(args []string, root string) bool { + if !slices.Contains(args, "diff") && !slices.Contains(args, "status") { + return false + } + ctx, cancel := context.WithTimeout(context.Background(), 3*time.Second) + defer cancel() + cmd := exec.CommandContext(ctx, "git", "config", "--includes", "--null", "--get-regexp", `^filter\..+\.(clean|smudge|process)$`) + cmd.Dir = root + cmd.Env = append(os.Environ(), "GIT_CONFIG_NOSYSTEM=1", "GIT_CONFIG_GLOBAL="+os.DevNull) + err := cmd.Run() + var exitError *exec.ExitError + switch { + case err == nil: + return true + case errors.As(err, &exitError): + return exitError.ExitCode() != 1 && exitError.ExitCode() != 128 + case errors.Is(err, exec.ErrNotFound): + return false + } + return true +} + +var ( + gitHardeningConfig = "-c core.fsmonitor=false -c log.showSignature=false" + gitDiffHardening = "--no-ext-diff --no-textconv" + gitDeniedPathspecs = `':(exclude,glob,icase)**/.env*' ':(exclude,glob,icase)**/*.pem' ` + + `':(exclude,glob,icase)**/*.key' ':(exclude,glob,icase)**/id_*'` +) + +// hardenSafeGit rewrites the git reads that run without approval so repo +// config cannot run programs (fsmonitor, signature checks, textconv and +// external diff drivers) and patches never include denied files. +func hardenSafeGit(script, root string) string { + root = resolveRoot(root) + file, err := syntax.NewParser(syntax.Variant(syntax.LangBash)).Parse(strings.NewReader(script), "") + if err != nil { + return script + } + type insertion struct { + offset uint + text string + } + var inserts []insertion + syntax.Walk(file, func(node syntax.Node) bool { + call, ok := node.(*syntax.CallExpr) + if !ok || len(call.Args) < 2 { + return true + } + argv := make([]string, len(call.Args)) + for index, word := range call.Args { + if argv[index], ok = literalWord(word); !ok { + return true + } + } + if argv[0] != "git" || !isKnownSafeCommand(argv, root) { + return true + } + inserts = append(inserts, insertion{call.Args[0].End().Offset(), " " + gitHardeningConfig}) + sub := 1 + for argv[sub] == "--no-pager" { + sub++ + } + if !slices.Contains(gitDiffCommands, argv[sub]) { + return true + } + inserts = append(inserts, insertion{call.Args[sub].End().Offset(), " " + gitDiffHardening}) + if slices.Contains(argv, "--follow") { + return true + } + start, _ := gitPathStart(argv[sub+1:], root) + separator := " --" + switch { + case start >= 0 && argv[sub+1+start] == "--": + separator = "" + case start >= 0: + inserts = append(inserts, insertion{call.Args[sub+1+start].Pos().Offset(), "-- "}) + separator = "" + } + inserts = append(inserts, insertion{call.Args[len(call.Args)-1].End().Offset(), separator + " " + gitDeniedPathspecs}) + return true + }) + slices.SortStableFunc(inserts, func(a, b insertion) int { return int(a.offset) - int(b.offset) }) + for index := len(inserts) - 1; index >= 0; index-- { + insert := inserts[index] + script = script[:insert.offset] + insert.text + script[insert.offset:] + } + return script +} + +func gitBranchIsReadOnly(args []string) bool { + if len(args) == 0 { + return true + } + for _, arg := range args { + if !slices.Contains([]string{"--list", "-l", "--show-current", "-a", "--all", "-r", "--remotes", "-v", "-vv", "--verbose"}, arg) && + !strings.HasPrefix(arg, "--format=") { + return false + } + } + return true +} + +// scriptForbidden matches forbidden rules against every command a script +// could run, whatever its redirects, expansions or nesting: each simple +// command in the AST, the command behind wrappers such as sudo or env, and +// bash -c / eval bodies. A script that doesn't parse is matched token by +// token. It returns the justification of the first forbidden match. +func scriptForbidden(policy *Policy, script string) string { + for _, argv := range forbiddenCandidates(script, 0) { + if decision, justification := policy.Evaluate(argv); decision == DecisionForbidden { + if justification == "" { + justification = strings.Join(argv, " ") + } + return justification + } + } + return "" +} + +func forbiddenCandidates(script string, depth int) [][]string { + file, err := syntax.NewParser(syntax.Variant(syntax.LangBash)).Parse(strings.NewReader(script), "") + if err != nil || depth > 3 { + return tokenSuffixes(strings.FieldsFunc(script, func(r rune) bool { + return unicode.IsSpace(r) || strings.ContainsRune(";|&()<>`'\"${}\\", r) + })) + } + var candidates [][]string + syntax.Walk(file, func(node syntax.Node) bool { + call, ok := node.(*syntax.CallExpr) + if !ok || len(call.Args) == 0 { + return true + } + argv := make([]string, len(call.Args)) + for index, word := range call.Args { + argv[index] = looseWord(word) + } + candidates = append(candidates, commandCandidates(argv, depth)...) + return true + }) + return candidates +} + +// looseWord renders the literal parts of a word; anything dynamic becomes a +// placeholder no rule matches. +func looseWord(word *syntax.Word) string { + var value strings.Builder + for _, part := range word.Parts { + switch part := part.(type) { + case *syntax.Lit: + value.WriteString(part.Value) + case *syntax.SglQuoted: + value.WriteString(part.Value) + case *syntax.DblQuoted: + for _, inner := range part.Parts { + if lit, ok := inner.(*syntax.Lit); ok { + value.WriteString(lit.Value) + } else { + value.WriteString("\x00") + } + } + default: + value.WriteString("\x00") + } + } + return value.String() +} + +var shellInterpreters = map[string]bool{"sh": true, "bash": true, "zsh": true, "dash": true, "ksh": true} + +func commandCandidates(argv []string, depth int) [][]string { + if len(argv) == 0 { + return nil + } + candidates := [][]string{argv} + name := filepath.Base(argv[0]) + if name != argv[0] { + candidates = append(candidates, append([]string{name}, argv[1:]...)) + } + switch { + case name == "eval": + candidates = append(candidates, forbiddenCandidates(strings.Join(argv[1:], " "), depth+1)...) + case shellInterpreters[name]: + for index, arg := range argv[1:] { + if strings.HasPrefix(arg, "-") && strings.Contains(arg, "c") && index+2 < len(argv) { + candidates = append(candidates, forbiddenCandidates(argv[index+2], depth+1)...) + break + } + } + case unrememberableCommands[name]: + for index := 1; index < len(argv); index++ { + candidates = append(candidates, commandCandidates(argv[index:], depth)...) + } + } + return candidates +} + +func tokenSuffixes(tokens []string) [][]string { + suffixes := make([][]string, 0, len(tokens)*2) + for index := range tokens { + suffixes = append(suffixes, tokens[index:]) + if base := filepath.Base(tokens[index]); base != tokens[index] { + suffixes = append(suffixes, append([]string{base}, tokens[index+1:]...)) + } + } + return suffixes +} diff --git a/internal/mcp/builtin_manager_test.go b/internal/mcp/builtin_manager_test.go new file mode 100644 index 0000000..af9854a --- /dev/null +++ b/internal/mcp/builtin_manager_test.go @@ -0,0 +1,306 @@ +package mcp + +import ( + "context" + "encoding/json" + "fmt" + "os" + "path/filepath" + "strings" + "testing" + + sdk "github.com/modelcontextprotocol/go-sdk/mcp" +) + +func TestBuiltinBashIsOptInAndStatePersists(t *testing.T) { + manager, options := newBuiltinManager(t, func(string) (string, error) { return "", fmt.Errorf("not installed") }) + tools, _ := manager.ClientTools() + if containsClientTool(tools, "local_bash") { + t.Fatal("bash must be off by default") + } + if result := manager.Call(context.Background(), "local_bash", map[string]any{"command": "printf unexpected"}); !result.IsError { + t.Fatal("disabled bash unexpectedly ran") + } + if err := manager.SetBuiltinEnabled("local_bash", true); err != nil { + t.Fatal(err) + } + tools, _ = manager.ClientTools() + if !containsClientTool(tools, "local_bash") { + t.Fatal("bash was not advertised after opt-in") + } + if !manager.NeedsApproval("local_bash", map[string]any{"command": "printf safe"}) { + t.Fatal("bash must always require approval") + } + result := manager.Call(context.Background(), "local_bash", map[string]any{"command": "printf opt-in"}) + if result.IsError || !strings.Contains(result.Output, "opt-in") { + t.Fatalf("opted-in bash did not run: %#v", result) + } + reloaded := NewManager(options) + defer reloaded.Close() + if err := reloaded.Load(); err != nil { + t.Fatal(err) + } + if !containsBuiltin(reloaded.BuiltinTools(), "local_bash", true) { + t.Fatal("bash enablement did not persist in MCP state") + } +} + +func TestCombinedClientToolLimitsAndBuiltinNameReservation(t *testing.T) { + manager, _ := newBuiltinManager(t, func(string) (string, error) { return "/bin/true", nil }) + addRuntimeTools(manager, "sentry", []*sdk.Tool{ + {Name: "cli", InputSchema: map[string]any{"type": "object"}, Annotations: &sdk.ToolAnnotations{ReadOnlyHint: true}}, + }) + mcpTools := make([]*sdk.Tool, 59) + for index := range mcpTools { + mcpTools[index] = &sdk.Tool{ + Name: fmt.Sprintf("tool_%02d", index), + Description: "MCP tool", + InputSchema: map[string]any{"type": "object"}, + } + } + addRuntimeTools(manager, "workspace", mcpTools) + + tools, omitted := manager.ClientTools() + if len(tools) > MaxClientTools || omitted == 0 { + t.Fatalf("combined client tool count = %d, omitted = %d", len(tools), omitted) + } + if !containsClientTool(tools, "local_sentry_cli") { + t.Fatal("built-in sentry-cli name is missing") + } + if !containsClientTool(tools, "local_sentry_cli_2") { + t.Fatal("MCP tool did not receive a collision-safe name") + } + if ref, ok := manager.Lookup("local_sentry_cli_2"); !ok || ref.ServerID == builtinServerID { + t.Fatalf("collision-safe MCP reference is invalid: %+v %t", ref, ok) + } + if manager.NeedsApproval("local_sentry_cli_2", nil) { + t.Fatal("read-only MCP approval behavior changed during name reservation") + } + for _, tool := range tools { + if len(tool.InputSchema) > MaxSchemaBytes { + t.Fatalf("%s schema exceeds limit: %d", tool.Name, len(tool.InputSchema)) + } + } + encoded, err := json.Marshal(tools) + if err != nil { + t.Fatal(err) + } + if len(encoded) > MaxClientToolsBytes { + t.Fatalf("combined schema payload is %d bytes, limit is %d", len(encoded), MaxClientToolsBytes) + } +} + +func TestCombinedSchemaByteLimitWithMCPTools(t *testing.T) { + manager, _ := newBuiltinManager(t, func(string) (string, error) { return "", fmt.Errorf("not installed") }) + tools := make([]*sdk.Tool, 40) + for index := range tools { + tools[index] = &sdk.Tool{ + Name: fmt.Sprintf("large_%02d", index), + Description: "large schema", + InputSchema: map[string]any{ + "type": "object", + "properties": map[string]any{ + "value": map[string]any{"type": "string", "description": strings.Repeat("x", 10_000)}, + }, + }, + } + } + tools[0].InputSchema = map[string]any{ + "type": "object", + "properties": map[string]any{ + "value": map[string]any{"type": "string", "description": strings.Repeat("x", MaxSchemaBytes+1)}, + }, + } + addRuntimeTools(manager, "large", tools) + + advertised, omitted := manager.ClientTools() + if omitted == 0 || len(advertised) >= MaxClientTools { + t.Fatalf("large schemas were not limited: %d tools, %d omitted", len(advertised), omitted) + } + for _, tool := range advertised { + if len(tool.InputSchema) > MaxSchemaBytes { + t.Fatalf("%s schema exceeds limit: %d", tool.Name, len(tool.InputSchema)) + } + } + encoded, err := json.Marshal(advertised) + if err != nil { + t.Fatal(err) + } + if len(encoded) > MaxClientToolsBytes { + t.Fatalf("combined schema payload is %d bytes, limit is %d", len(encoded), MaxClientToolsBytes) + } +} + +func newBuiltinManager(t *testing.T, lookPath func(string) (string, error)) (*Manager, ManagerOptions) { + t.Helper() + cwd, home := t.TempDir(), t.TempDir() + options := ManagerOptions{ + CWD: cwd, + Store: &Store{UserPath: filepath.Join(home, "mcp.json"), CWD: cwd}, + State: &StateStore{Path: filepath.Join(home, "mcp-state.json")}, + Discovery: &Discovery{Home: home, CWD: cwd}, + LookPath: lookPath, + } + manager := NewManager(options) + if err := manager.Load(); err != nil { + t.Fatal(err) + } + t.Cleanup(manager.Close) + return manager, options +} + +func addRuntimeTools(manager *Manager, name string, tools []*sdk.Tool) { + server := Server{Name: name, Scope: ScopeUser, Path: filepath.Join("project", name+".mcp.json")} + runtime := &runtimeServer{server: server, status: StatusConnected, tools: tools} + manager.mu.Lock() + manager.servers = append(manager.servers, runtime) + manager.byID[server.ID()] = runtime + manager.mu.Unlock() +} + +func containsClientTool(tools []ClientTool, name string) bool { + for _, tool := range tools { + if tool.Name == name { + return true + } + } + return false +} + +func containsBuiltin(tools []BuiltinToolInfo, name string, enabled bool) bool { + for _, tool := range tools { + if tool.Name == name { + return tool.Enabled == enabled + } + } + return false +} + +func TestManagerRemembersApprovalsAndReportsBadRules(t *testing.T) { + home := t.TempDir() + t.Setenv("HOME", home) + rules := filepath.Join(home, ".rootly-cli", "rules") + if err := os.MkdirAll(rules, 0o700); err != nil { + t.Fatal(err) + } + if err := os.WriteFile(filepath.Join(rules, "broken.rules"), []byte(`prefix_rule(`), 0o600); err != nil { + t.Fatal(err) + } + cwd := t.TempDir() + options := ManagerOptions{ + CWD: cwd, + Store: &Store{UserPath: filepath.Join(home, "mcp.json"), CWD: cwd}, + State: &StateStore{Path: filepath.Join(home, "mcp-state.json")}, + Discovery: &Discovery{Home: home, CWD: cwd}, + LookPath: func(string) (string, error) { return "/bin/true", nil }, + } + manager := NewManager(options) + defer manager.Close() + if err := manager.Load(); err == nil || !strings.Contains(err.Error(), "broken.rules:1:13") { + t.Fatalf("bad rules file not reported: %v", err) + } + if err := manager.SetBuiltinEnabled("local_bash", true); err != nil { + t.Fatal(err) + } + manager.ClientTools() + plan := map[string]any{"command": "terraform plan -out=tf.plan"} + if !manager.NeedsApproval("local_bash", plan) { + t.Fatal("terraform plan should ask before it is remembered") + } + if got := manager.ApprovalOptions("local_bash", plan); !got.Persist || strings.Join(got.Prefix, " ") != "terraform plan" { + t.Fatalf("options = %+v", got) + } + if err := manager.RememberApproval("local_bash", plan, false); err != nil { + t.Fatal(err) + } + if manager.NeedsApproval("local_bash", map[string]any{"command": "terraform plan"}) { + t.Fatal("session approval not applied") + } + complexCommand := map[string]any{"command": "terraform plan > $HOME/plan"} + if got := manager.ApprovalOptions("local_bash", complexCommand); got.Session || got.Persist { + t.Fatalf("complex shell must be once-only, got %+v", got) + } + if err := manager.RememberApproval("local_bash", complexCommand, true); err == nil { + t.Fatal("complex shell must not be saved") + } + if err := manager.RememberApproval("local_bash", map[string]any{"command": "npm test"}, true); err != nil { + t.Fatal(err) + } + if err := os.Remove(filepath.Join(rules, "broken.rules")); err != nil { + t.Fatal(err) + } + fresh := NewManager(options) + defer fresh.Close() + if err := fresh.Load(); err != nil { + t.Fatal(err) + } + if err := fresh.SetBuiltinEnabled("local_bash", true); err != nil { + t.Fatal(err) + } + fresh.ClientTools() + if fresh.NeedsApproval("local_bash", map[string]any{"command": "npm test -- --watch=false"}) { + t.Fatal("saved rule not applied in a new process") + } + if !fresh.NeedsApproval("local_bash", plan) { + t.Fatal("session approval leaked into a new process") + } + + addRuntimeTools(manager, "deploys", []*sdk.Tool{{Name: "rollback", InputSchema: map[string]any{"type": "object"}}}) + all, _ := manager.ClientTools() + name := "" + for _, tool := range all { + if strings.HasSuffix(tool.Name, "rollback") { + name = tool.Name + } + } + if got := manager.ApprovalOptions(name, nil); !got.Session || got.Persist { + t.Fatalf("MCP tool options = %+v", got) + } + if err := manager.RememberApproval(name, nil, true); err == nil { + t.Fatal("MCP tools cannot be saved as rules") + } + if err := manager.RememberApproval(name, nil, false); err != nil || manager.NeedsApproval(name, nil) { + t.Fatalf("MCP session approval: %v", err) + } +} + +func TestSessionApprovalStaysWithItsServer(t *testing.T) { + manager := NewManager(ManagerOptions{Store: NewStore(t.TempDir())}) + write := []*sdk.Tool{{Name: "write", InputSchema: map[string]any{"type": "object"}}} + user := &runtimeServer{server: Server{Name: "docs", Scope: ScopeUser, Path: "user.json"}, status: StatusConnected, tools: write} + project := &runtimeServer{server: Server{Name: "docs", Scope: ScopeProject, Path: "project.json"}, status: StatusConnected, tools: write} + manager.mu.Lock() + manager.servers = append(manager.servers, user, project) + manager.byID[user.server.ID()] = user + manager.byID[project.server.ID()] = project + manager.mu.Unlock() + nameFor := func(runtime *runtimeServer) string { + tools, _ := manager.ClientTools() + for _, tool := range tools { + if ref, ok := manager.Lookup(tool.Name); ok && ref.ServerID == runtime.server.ID() && ref.Tool == "write" { + return tool.Name + } + } + t.Fatalf("no write tool for %s", runtime.server.ID()) + return "" + } + approved := nameFor(user) + if err := manager.RememberApproval(approved, nil, false); err != nil || manager.NeedsApproval(approved, nil) { + t.Fatalf("session approval for the user server: %v", err) + } + if other := nameFor(project); !manager.NeedsApproval(other, nil) { + t.Fatal("the other docs server inherited the approval") + } + manager.mu.Lock() + user.sessionOff = true + manager.mu.Unlock() + if other := nameFor(project); !manager.NeedsApproval(other, nil) { + t.Fatalf("%s moved to the other server and kept the approval", other) + } + manager.mu.Lock() + user.sessionOff = false + manager.mu.Unlock() + if manager.NeedsApproval(nameFor(user), nil) { + t.Fatal("approval lost for the server it was given to") + } +} diff --git a/internal/mcp/config.go b/internal/mcp/config.go new file mode 100644 index 0000000..7a7f916 --- /dev/null +++ b/internal/mcp/config.go @@ -0,0 +1,410 @@ +package mcp + +import ( + "crypto/sha256" + "encoding/hex" + "encoding/json" + "errors" + "fmt" + "os" + "path/filepath" + "regexp" + "sort" + "strings" + "time" + + "github.com/rootlyhq/rootly-cli/internal/config" +) + +const ( + TransportStdio = "stdio" + TransportHTTP = "http" + TransportSSE = "sse" + + ScopeLocal = "local" + ScopeUser = "user" + ScopeProject = "project" + ScopeImported = "imported" + ScopeBuiltin = "builtin" + + userConfigFile = "mcp.json" + projectConfigFile = ".mcp.json" +) + +var serverNamePattern = regexp.MustCompile(`^[A-Za-z0-9][A-Za-z0-9_.-]{0,63}$`) + +type ServerConfig struct { + Type string `json:"type,omitempty"` + Command string `json:"command,omitempty"` + Args []string `json:"args,omitempty"` + Env map[string]string `json:"env,omitempty"` + URL string `json:"url,omitempty"` + Headers map[string]string `json:"headers,omitempty"` + BearerTokenEnvVar string `json:"bearer_token_env_var,omitempty"` + Timeout int `json:"timeout,omitempty"` +} + +// MaxTimeoutSeconds bounds the optional per-server timeout. +const MaxTimeoutSeconds = 600 + +// CallTimeout returns the per-server timeout, or fallback when unset. +func (c ServerConfig) CallTimeout(fallback time.Duration) time.Duration { + if c.Timeout <= 0 { + return fallback + } + return time.Duration(c.Timeout) * time.Second +} + +func (c ServerConfig) Transport() string { + switch strings.ToLower(c.Type) { + case TransportHTTP, "streamable-http", "streamable_http", "streamablehttp": + return TransportHTTP + case TransportSSE: + return TransportSSE + case "", TransportStdio: + if c.Command == "" && c.URL != "" { + return TransportHTTP + } + return TransportStdio + default: + return strings.ToLower(c.Type) + } +} + +func (c ServerConfig) Validate() error { + if c.Timeout < 0 || c.Timeout > MaxTimeoutSeconds { + return fmt.Errorf("timeout must be between 1 and %d seconds", MaxTimeoutSeconds) + } + switch c.Transport() { + case TransportStdio: + if strings.TrimSpace(c.Command) == "" { + return errors.New("stdio server needs a command") + } + if c.URL != "" { + return errors.New("stdio server can't have a url; set type to http or sse") + } + case TransportHTTP, TransportSSE: + if strings.TrimSpace(c.URL) == "" { + return fmt.Errorf("%s server needs a url", c.Transport()) + } + if c.Command != "" { + return fmt.Errorf("%s server can't have a command", c.Transport()) + } + default: + return fmt.Errorf("unsupported transport %q (expected stdio, http or sse)", c.Type) + } + return nil +} + +func (c ServerConfig) Hash() string { + data, _ := json.Marshal(c) + sum := sha256.Sum256(data) + return hex.EncodeToString(sum[:]) +} + +type Server struct { + Name string + Config ServerConfig + Scope string + Path string + SourceLabel string + // Project is the project directory when an imported server comes from a + // project-controlled file, so it needs approval like a .mcp.json server. + Project string +} + +// NeedsProjectTrust reports whether the server's config is controlled by the +// project checkout and must be approved (per config hash) before it starts. +func (s Server) NeedsProjectTrust() bool { + return s.Scope == ScopeProject || s.Project != "" +} + +func (s Server) approvalPath() string { + if s.Scope == ScopeImported && s.Project != "" { + return s.Project + "#" + s.Path + } + return s.Path +} + +func (s Server) ID() string { + return s.Scope + ":" + s.Path + ":" + s.Name +} + +func (s Server) ReadOnly() bool { + return s.Scope == ScopeImported +} + +func ValidateServerName(name string) error { + if !serverNamePattern.MatchString(name) { + return fmt.Errorf("invalid server name %q: use letters, digits, '-', '_' or '.' (max 64)", name) + } + return nil +} + +func UserConfigPath() string { + return filepath.Join(config.Dir(), userConfigFile) +} + +func ProjectConfigPath(cwd string) string { + return filepath.Join(cwd, projectConfigFile) +} + +type jsonDocument map[string]json.RawMessage + +func readJSONDocument(path string) (jsonDocument, error) { + data, err := os.ReadFile(path) + if err != nil { + if os.IsNotExist(err) { + return jsonDocument{}, nil + } + return nil, err + } + if strings.TrimSpace(string(data)) == "" { + return jsonDocument{}, nil + } + document := jsonDocument{} + if err := json.Unmarshal(data, &document); err != nil { + return nil, fmt.Errorf("parse %s: %w", path, err) + } + return document, nil +} + +func writeJSONDocument(path string, document jsonDocument, perm os.FileMode) error { + if err := os.MkdirAll(filepath.Dir(path), 0o700); err != nil { + return err + } + data, err := json.MarshalIndent(document, "", " ") + if err != nil { + return err + } + data = append(data, '\n') + tmp, err := os.CreateTemp(filepath.Dir(path), "."+filepath.Base(path)+".*") + if err != nil { + return err + } + defer func() { _ = os.Remove(tmp.Name()) }() + if _, err := tmp.Write(data); err != nil { + _ = tmp.Close() + return err + } + if err := tmp.Close(); err != nil { + return err + } + if info, statErr := os.Stat(path); statErr == nil { + perm = info.Mode().Perm() + } + if err := os.Chmod(tmp.Name(), perm); err != nil { + return err + } + return os.Rename(tmp.Name(), path) +} + +func decodeServers(raw json.RawMessage) (map[string]ServerConfig, error) { + if len(raw) == 0 || string(raw) == "null" { + return map[string]ServerConfig{}, nil + } + servers := map[string]ServerConfig{} + if err := json.Unmarshal(raw, &servers); err != nil { + return nil, err + } + return servers, nil +} + +func decodeServersLenient(raw json.RawMessage) map[string]ServerConfig { + if len(raw) == 0 { + return map[string]ServerConfig{} + } + entries := map[string]json.RawMessage{} + if json.Unmarshal(raw, &entries) != nil { + return map[string]ServerConfig{} + } + servers := make(map[string]ServerConfig, len(entries)) + for name, entry := range entries { + var server ServerConfig + if json.Unmarshal(entry, &server) == nil { + servers[name] = server + } + } + return servers +} + +func sortedServers(servers map[string]ServerConfig, scope, path, label string) []Server { + names := make([]string, 0, len(servers)) + for name := range servers { + names = append(names, name) + } + sort.Strings(names) + result := make([]Server, 0, len(names)) + for _, name := range names { + result = append(result, Server{Name: name, Config: servers[name], Scope: scope, Path: path, SourceLabel: label}) + } + return result +} + +// Store reads and writes the scopes rootly owns: user and local in +// ~/.rootly-cli/mcp.json, project in ./.mcp.json. +type Store struct { + UserPath string + CWD string +} + +func NewStore(cwd string) *Store { + return &Store{UserPath: UserConfigPath(), CWD: cwd} +} + +func (s *Store) ProjectPath() string { + return ProjectConfigPath(s.CWD) +} + +func (s *Store) Servers(scope string) ([]Server, error) { + switch scope { + case ScopeUser: + document, err := readJSONDocument(s.UserPath) + if err != nil { + return nil, err + } + servers, err := decodeServers(document["mcpServers"]) + if err != nil { + return nil, fmt.Errorf("parse %s mcpServers: %w", s.UserPath, err) + } + return sortedServers(servers, ScopeUser, s.UserPath, "rootly"), nil + case ScopeLocal: + document, err := readJSONDocument(s.UserPath) + if err != nil { + return nil, err + } + projects := map[string]jsonDocument{} + if raw, ok := document["projects"]; ok { + if err := json.Unmarshal(raw, &projects); err != nil { + return nil, fmt.Errorf("parse %s projects: %w", s.UserPath, err) + } + } + servers, err := decodeServers(projects[s.CWD]["mcpServers"]) + if err != nil { + return nil, fmt.Errorf("parse %s projects[%q]: %w", s.UserPath, s.CWD, err) + } + return sortedServers(servers, ScopeLocal, s.UserPath, "rootly"), nil + case ScopeProject: + path := s.ProjectPath() + document, err := readJSONDocument(path) + if err != nil { + return nil, err + } + servers, err := decodeServers(document["mcpServers"]) + if err != nil { + return nil, fmt.Errorf("parse %s mcpServers: %w", path, err) + } + return sortedServers(servers, ScopeProject, path, "project"), nil + default: + return nil, fmt.Errorf("unknown scope %q (expected local, user or project)", scope) + } +} + +// All returns every writable-scope server, local first, then project, then user. +func (s *Store) All() ([]Server, error) { + var all []Server + var errs []error + for _, scope := range []string{ScopeLocal, ScopeProject, ScopeUser} { + servers, err := s.Servers(scope) + if err != nil { + errs = append(errs, err) + continue + } + all = append(all, servers...) + } + return all, errors.Join(errs...) +} + +func (s *Store) Add(scope, name string, server ServerConfig) error { + if err := ValidateServerName(name); err != nil { + return err + } + if err := server.Validate(); err != nil { + return err + } + return s.update(scope, func(servers map[string]ServerConfig) error { + if _, exists := servers[name]; exists { + return fmt.Errorf("MCP server %q already exists in %s config", name, scope) + } + servers[name] = server + return nil + }) +} + +func (s *Store) Remove(scope, name string) error { + return s.update(scope, func(servers map[string]ServerConfig) error { + if _, exists := servers[name]; !exists { + return fmt.Errorf("no MCP server named %q in %s config", name, scope) + } + delete(servers, name) + return nil + }) +} + +func (s *Store) update(scope string, change func(map[string]ServerConfig) error) error { + switch scope { + case ScopeUser, ScopeProject: + path := s.UserPath + perm := os.FileMode(0o600) + if scope == ScopeProject { + path, perm = s.ProjectPath(), 0o644 + } + document, err := readJSONDocument(path) + if err != nil { + return err + } + servers, err := decodeServers(document["mcpServers"]) + if err != nil { + return fmt.Errorf("parse %s mcpServers: %w", path, err) + } + if err := change(servers); err != nil { + return err + } + encoded, err := json.Marshal(servers) + if err != nil { + return err + } + document["mcpServers"] = encoded + return writeJSONDocument(path, document, perm) + case ScopeLocal: + document, err := readJSONDocument(s.UserPath) + if err != nil { + return err + } + projects := map[string]jsonDocument{} + if raw, ok := document["projects"]; ok { + if err := json.Unmarshal(raw, &projects); err != nil { + return fmt.Errorf("parse %s projects: %w", s.UserPath, err) + } + } + project := projects[s.CWD] + if project == nil { + project = jsonDocument{} + } + servers, err := decodeServers(project["mcpServers"]) + if err != nil { + return err + } + if err := change(servers); err != nil { + return err + } + encoded, err := json.Marshal(servers) + if err != nil { + return err + } + project["mcpServers"] = encoded + if len(servers) == 0 && len(project) == 1 { + delete(projects, s.CWD) + } else { + projects[s.CWD] = project + } + encodedProjects, err := json.Marshal(projects) + if err != nil { + return err + } + document["projects"] = encodedProjects + return writeJSONDocument(s.UserPath, document, 0o600) + default: + return fmt.Errorf("unknown scope %q (expected local, user or project)", scope) + } +} diff --git a/internal/mcp/expand.go b/internal/mcp/expand.go new file mode 100644 index 0000000..017da0c --- /dev/null +++ b/internal/mcp/expand.go @@ -0,0 +1,48 @@ +package mcp + +import ( + "os" + "regexp" +) + +var envReference = regexp.MustCompile(`\$\{([A-Za-z_][A-Za-z0-9_]*)(:-([^}]*))?\}`) + +// ExpandEnv replaces ${VAR} and ${VAR:-default}; an unset ${VAR} becomes "". +func ExpandEnv(value string, lookup func(string) (string, bool)) string { + if lookup == nil { + lookup = os.LookupEnv + } + return envReference.ReplaceAllStringFunc(value, func(match string) string { + parts := envReference.FindStringSubmatch(match) + if current, ok := lookup(parts[1]); ok && (current != "" || parts[2] == "") { + return current + } + return parts[3] + }) +} + +func (c ServerConfig) Expanded(lookup func(string) (string, bool)) ServerConfig { + expanded := ServerConfig{ + Type: c.Type, + Command: ExpandEnv(c.Command, lookup), + URL: ExpandEnv(c.URL, lookup), + BearerTokenEnvVar: c.BearerTokenEnvVar, + Timeout: c.Timeout, + } + for _, arg := range c.Args { + expanded.Args = append(expanded.Args, ExpandEnv(arg, lookup)) + } + if len(c.Env) > 0 { + expanded.Env = make(map[string]string, len(c.Env)) + for key, value := range c.Env { + expanded.Env[key] = ExpandEnv(value, lookup) + } + } + if len(c.Headers) > 0 { + expanded.Headers = make(map[string]string, len(c.Headers)) + for key, value := range c.Headers { + expanded.Headers[key] = ExpandEnv(value, lookup) + } + } + return expanded +} diff --git a/internal/mcp/freshness_test.go b/internal/mcp/freshness_test.go new file mode 100644 index 0000000..71fa35c --- /dev/null +++ b/internal/mcp/freshness_test.go @@ -0,0 +1,228 @@ +package mcp + +import ( + "context" + "encoding/json" + "errors" + "io" + "net/http" + "net/http/httptest" + "path/filepath" + "strings" + "sync/atomic" + "testing" + "time" +) + +func TestServerTimeoutConfig(t *testing.T) { + config := ServerConfig{Command: "srv"} + if got := config.CallTimeout(45 * time.Second); got != 45*time.Second { + t.Errorf("unset timeout = %s, want fallback", got) + } + config.Timeout = 5 + if got := config.CallTimeout(45 * time.Second); got != 5*time.Second { + t.Errorf("timeout = %s, want 5s", got) + } + if err := config.Validate(); err != nil { + t.Errorf("valid timeout rejected: %v", err) + } + for _, bad := range []int{-1, MaxTimeoutSeconds + 1} { + config.Timeout = bad + if err := config.Validate(); err == nil { + t.Errorf("timeout %d accepted", bad) + } + } + var parsed ServerConfig + if err := json.Unmarshal([]byte(`{"command":"srv","timeout":30}`), &parsed); err != nil || parsed.Timeout != 30 { + t.Errorf("timeout not parsed from config: %+v %v", parsed, err) + } +} + +func refreshFixture(t *testing.T, refreshToken string) (*httptest.Server, *TokenStore, *atomic.Int32) { + t.Helper() + var refreshes atomic.Int32 + mux := http.NewServeMux() + mux.HandleFunc("/token", func(w http.ResponseWriter, r *http.Request) { + refreshes.Add(1) + w.Header().Set("Content-Type", "application/json") + _, _ = io.WriteString(w, `{"access_token":"fresh","refresh_token":"r2","token_type":"Bearer","expires_in":3600}`) + }) + mux.HandleFunc("/mcp", func(w http.ResponseWriter, r *http.Request) { + if r.Header.Get("Authorization") != "Bearer fresh" { + w.WriteHeader(http.StatusUnauthorized) + return + } + body, _ := io.ReadAll(r.Body) + _, _ = w.Write(body) + }) + server := httptest.NewServer(mux) + t.Cleanup(server.Close) + store := &TokenStore{Path: filepath.Join(t.TempDir(), "mcp-oauth.json")} + credential := StoredCredential{ClientID: "c", TokenURL: server.URL + "/token", AccessToken: "stale", RefreshToken: refreshToken, TokenType: "Bearer", Expiry: time.Now().Add(time.Hour)} + if err := store.Put(server.URL+"/mcp", credential); err != nil { + t.Fatal(err) + } + return server, store, &refreshes +} + +func TestRemote401RefreshesTokenAndRetriesOnce(t *testing.T) { + server, store, refreshes := refreshFixture(t, "r1") + manager := &Manager{options: ManagerOptions{Tokens: store}} + config := ServerConfig{Type: TransportHTTP, URL: server.URL + "/mcp"} + bearer, refresh := manager.bearerFuncs(config) + probe := &authProbe{} + client := newServerHTTPClientWithRefresh(config, bearer, refresh, probe) + + request, _ := http.NewRequestWithContext(t.Context(), http.MethodPost, config.URL, strings.NewReader(`{"ping":1}`)) + response, err := client.Do(request) + if err != nil { + t.Fatal(err) + } + body, _ := io.ReadAll(response.Body) + _ = response.Body.Close() + if response.StatusCode != http.StatusOK || string(body) != `{"ping":1}` { + t.Fatalf("retry = %d %q, want 200 with replayed body", response.StatusCode, body) + } + if refreshes.Load() != 1 || probe.seen() { + t.Errorf("refreshes = %d, needs-login = %v; want one silent refresh", refreshes.Load(), probe.seen()) + } + if stored, _ := store.Get(config.URL); stored.AccessToken != "fresh" { + t.Errorf("refreshed token not persisted: %q", stored.AccessToken) + } +} + +func TestRemote401WithoutRefreshTokenNeedsLogin(t *testing.T) { + server, store, refreshes := refreshFixture(t, "") + manager := &Manager{options: ManagerOptions{Tokens: store}} + config := ServerConfig{Type: TransportHTTP, URL: server.URL + "/mcp"} + bearer, refresh := manager.bearerFuncs(config) + probe := &authProbe{} + client := newServerHTTPClientWithRefresh(config, bearer, refresh, probe) + + request, _ := http.NewRequestWithContext(t.Context(), http.MethodGet, config.URL, http.NoBody) + response, err := client.Do(request) + if err != nil { + t.Fatal(err) + } + _ = response.Body.Close() + if response.StatusCode != http.StatusUnauthorized || !probe.seen() || refreshes.Load() != 0 { + t.Errorf("status %d, needs-login %v, refreshes %d; want 401 + needs-login without refresh", response.StatusCode, probe.seen(), refreshes.Load()) + } +} + +func TestExpandedKeepsTimeout(t *testing.T) { + config := ServerConfig{Command: "x", Timeout: 5}.Expanded(func(string) (string, bool) { return "", false }) + if config.CallTimeout(time.Minute) != 5*time.Second { + t.Fatalf("expanded timeout = %v, want 5s", config.CallTimeout(time.Minute)) + } +} + +func TestRemote401RefreshStopsWithRequestContext(t *testing.T) { + release := make(chan struct{}) + mux := http.NewServeMux() + mux.HandleFunc("/token", func(w http.ResponseWriter, r *http.Request) { + _, _ = io.ReadAll(r.Body) + select { + case <-r.Context().Done(): + case <-release: + } + }) + mux.HandleFunc("/mcp", func(w http.ResponseWriter, r *http.Request) { + w.WriteHeader(http.StatusUnauthorized) + }) + server := httptest.NewServer(mux) + t.Cleanup(server.Close) + t.Cleanup(func() { close(release) }) + store := &TokenStore{Path: filepath.Join(t.TempDir(), "mcp-oauth.json")} + if err := store.Put(server.URL+"/mcp", StoredCredential{ClientID: "c", TokenURL: server.URL + "/token", AccessToken: "stale", RefreshToken: "r1", TokenType: "Bearer", Expiry: time.Now().Add(time.Hour)}); err != nil { + t.Fatal(err) + } + manager := &Manager{options: ManagerOptions{Tokens: store}} + config := ServerConfig{Type: TransportHTTP, URL: server.URL + "/mcp"} + bearer, refresh := manager.bearerFuncs(config) + client := newServerHTTPClientWithRefresh(config, bearer, refresh, &authProbe{}) + + ctx, cancel := context.WithTimeout(t.Context(), 200*time.Millisecond) + defer cancel() + request, _ := http.NewRequestWithContext(ctx, http.MethodGet, config.URL, http.NoBody) + started := time.Now() + response, err := client.Do(request) + if err == nil { + _ = response.Body.Close() + } + if elapsed := time.Since(started); elapsed > 3*time.Second { + t.Fatalf("refresh ignored request cancellation, took %v", elapsed) + } +} + +func TestExpiredBearerRefreshStopsWithRequestContext(t *testing.T) { + release := make(chan struct{}) + tokenStarted := make(chan struct{}, 1) + var mcpRequests atomic.Int32 + mux := http.NewServeMux() + mux.HandleFunc("/token", func(w http.ResponseWriter, r *http.Request) { + _, _ = io.ReadAll(r.Body) + tokenStarted <- struct{}{} + select { + case <-r.Context().Done(): + case <-release: + } + }) + mux.HandleFunc("/mcp", func(w http.ResponseWriter, r *http.Request) { + mcpRequests.Add(1) + }) + server := httptest.NewServer(mux) + t.Cleanup(server.Close) + t.Cleanup(func() { close(release) }) + store := &TokenStore{Path: filepath.Join(t.TempDir(), "mcp-oauth.json")} + if err := store.Put(server.URL+"/mcp", StoredCredential{ + ClientID: "c", + TokenURL: server.URL + "/token", + AccessToken: "expired", + RefreshToken: "r1", + TokenType: "Bearer", + Expiry: time.Now().Add(-time.Hour), + }); err != nil { + t.Fatal(err) + } + manager := &Manager{options: ManagerOptions{Tokens: store}} + config := ServerConfig{Type: TransportHTTP, URL: server.URL + "/mcp"} + bearer, refresh := manager.bearerFuncs(config) + client := newServerHTTPClientWithRefresh(config, bearer, refresh, &authProbe{}) + + ctx, cancel := context.WithTimeout(t.Context(), 200*time.Millisecond) + defer cancel() + request, _ := http.NewRequestWithContext(ctx, http.MethodGet, config.URL, http.NoBody) + started := time.Now() + response, err := client.Do(request) + if response != nil { + _ = response.Body.Close() + } + if !errors.Is(err, context.DeadlineExceeded) { + t.Fatalf("request error = %v, want context deadline exceeded", err) + } + if elapsed := time.Since(started); elapsed > 2*time.Second { + t.Fatalf("expired bearer refresh ignored request cancellation, took %v", elapsed) + } + select { + case <-tokenStarted: + default: + t.Fatal("expired token did not attempt refresh before the MCP request") + } + if got := mcpRequests.Load(); got != 0 { + t.Errorf("MCP requests = %d, want none after the token refresh timed out", got) + } +} + +func TestConnectUsesConfiguredTimeout(t *testing.T) { + manager := NewManager(ManagerOptions{CWD: t.TempDir()}) + server := Server{Name: "slow", Config: ServerConfig{Command: "sleep", Args: []string{"30"}, Timeout: 1}} + start := time.Now() + session, _, status, _, _ := manager.connect(context.Background(), server) + if session != nil || status == StatusConnected { + t.Fatalf("hanging server connected: %s", status) + } + if elapsed := time.Since(start); elapsed > 10*time.Second { + t.Fatalf("connect took %s; the 1s server timeout was ignored", elapsed) + } +} diff --git a/internal/mcp/lifecycle_test.go b/internal/mcp/lifecycle_test.go new file mode 100644 index 0000000..59cf784 --- /dev/null +++ b/internal/mcp/lifecycle_test.go @@ -0,0 +1,220 @@ +package mcp + +import ( + "context" + "encoding/json" + "net/http" + "net/http/httptest" + "os" + "path/filepath" + "strings" + "testing" + "time" + + sdk "github.com/modelcontextprotocol/go-sdk/mcp" +) + +const stdioHelperEnv = "ROOTLY_MCP_TEST_STDIO_SERVER" + +// TestMain lets the test binary act as a real stdio MCP server so the +// manager lifecycle runs against an actual child process. +func TestMain(m *testing.M) { + if os.Getenv(stdioHelperEnv) == "1" { + server := sdk.NewServer(&sdk.Implementation{Name: "helper", Version: "1"}, nil) + server.AddTool(&sdk.Tool{Name: "echo", Description: "Echo text", InputSchema: json.RawMessage(`{"type":"object","properties":{"text":{"type":"string"}}}`)}, + func(_ context.Context, request *sdk.CallToolRequest) (*sdk.CallToolResult, error) { + var args struct{ Text string } + _ = json.Unmarshal(request.Params.Arguments, &args) + if args.Text == "fail" { + return &sdk.CallToolResult{IsError: true, Content: []sdk.Content{&sdk.TextContent{Text: "tool failed"}}}, nil + } + return &sdk.CallToolResult{Content: []sdk.Content{&sdk.TextContent{Text: "echo: " + args.Text}}}, nil + }) + server.AddTool(&sdk.Tool{Name: "structured", Description: "Structured only", InputSchema: json.RawMessage(`{"type":"object"}`)}, + func(context.Context, *sdk.CallToolRequest) (*sdk.CallToolResult, error) { + return &sdk.CallToolResult{Content: []sdk.Content{}, StructuredContent: map[string]any{"ok": true}}, nil + }) + _ = server.Run(context.Background(), &sdk.StdioTransport{}) + os.Exit(0) + } + os.Exit(m.Run()) +} + +func stdioManager(t *testing.T, servers map[string]ServerConfig) *Manager { + t.Helper() + home := t.TempDir() + cwd := t.TempDir() + userPath := filepath.Join(home, ".rootly-cli", "mcp.json") + writeFile(t, userPath, marshalTestJSON(t, map[string]any{"mcpServers": servers})) + manager := NewManager(ManagerOptions{ + CWD: cwd, + Store: &Store{UserPath: userPath, CWD: cwd}, + State: &StateStore{Path: filepath.Join(home, ".rootly-cli", "mcp-state.json")}, + Tokens: &TokenStore{Path: filepath.Join(home, ".rootly-cli", "mcp-tokens.json")}, + Discovery: &Discovery{Home: home, CWD: cwd}, + LookPath: func(string) (string, error) { return "", os.ErrNotExist }, + }) + t.Cleanup(manager.Close) + if err := manager.Load(); err != nil { + t.Fatal(err) + } + return manager +} + +func waitConnected(t *testing.T, manager *Manager) { + t.Helper() + deadline := time.Now().Add(20 * time.Second) + for manager.Connecting() > 0 { + if time.Now().After(deadline) { + t.Fatalf("servers still connecting: %+v", manager.Views()) + } + select { + case <-manager.Changes(): + case <-time.After(50 * time.Millisecond): + } + } +} + +func viewByName(t *testing.T, manager *Manager, name string) ServerView { + t.Helper() + for _, view := range manager.Views() { + if view.Server.Name == name { + return view + } + } + t.Fatalf("server %q missing from %+v", name, manager.Views()) + return ServerView{} +} + +func toolNamed(t *testing.T, manager *Manager, suffix string) string { + t.Helper() + tools, _ := manager.ClientTools() + for _, tool := range tools { + if strings.HasSuffix(tool.Name, suffix) { + return tool.Name + } + } + t.Fatalf("no client tool ending in %q: %+v", suffix, tools) + return "" +} + +func TestManagerStdioLifecycle(t *testing.T) { + manager := stdioManager(t, map[string]ServerConfig{ + "helper": {Command: os.Args[0], Env: map[string]string{stdioHelperEnv: "1"}}, + "broken": {Command: filepath.Join(t.TempDir(), "missing-binary")}, + }) + manager.Start() + waitConnected(t, manager) + + helper := viewByName(t, manager, "helper") + if helper.Status != StatusConnected || len(helper.Tools) != 2 { + t.Fatalf("helper view = %+v", helper) + } + if broken := viewByName(t, manager, "broken"); broken.Status == StatusConnected || broken.Detail == "" { + t.Fatalf("missing binary should fail with a reason: %+v", broken) + } + + echo := toolNamed(t, manager, "echo") + if result := manager.Call(context.Background(), echo, map[string]any{"text": "hi"}); result.IsError || result.Output != "echo: hi" { + t.Fatalf("echo call = %+v", result) + } + if result := manager.Call(context.Background(), echo, map[string]any{"text": "fail"}); !result.IsError || result.Output != "tool failed" { + t.Fatalf("tool error not surfaced: %+v", result) + } + if result := manager.Call(context.Background(), toolNamed(t, manager, "structured"), nil); result.IsError || result.Output != `{"ok":true}` { + t.Fatalf("structured-only result = %+v", result) + } + if result := manager.Call(context.Background(), "nope_tool", nil); !result.IsError { + t.Fatalf("unknown tool must error: %+v", result) + } + + id := helper.Server.ID() + if err := manager.ToggleSession(id); err != nil { + t.Fatal(err) + } + if view := viewByName(t, manager, "helper"); !view.SessionOff { + t.Fatalf("toggle should switch the server off for this session: %+v", view) + } + for _, tool := range func() []ClientTool { tools, _ := manager.ClientTools(); return tools }() { + if strings.HasSuffix(tool.Name, "echo") { + t.Fatal("tools of a session-off server must not be advertised") + } + } + if result := manager.Call(context.Background(), echo, map[string]any{"text": "hi"}); !result.IsError { + t.Fatalf("calling a session-off server must fail: %+v", result) + } + if err := manager.ToggleSession(id); err != nil { + t.Fatal(err) + } + waitConnected(t, manager) + if err := manager.Reconnect(id); err != nil { + t.Fatal(err) + } + waitConnected(t, manager) + if view := viewByName(t, manager, "helper"); view.Status != StatusConnected || view.SessionOff { + t.Fatalf("reconnect should bring the server back: %+v", view) + } + if result := manager.Call(context.Background(), toolNamed(t, manager, "echo"), map[string]any{"text": "again"}); result.Output != "echo: again" { + t.Fatalf("call after reconnect = %+v", result) + } + if err := manager.ToggleSession("unknown"); err == nil { + t.Fatal("unknown server id must error") + } + if err := manager.Reconnect("unknown"); err == nil { + t.Fatal("unknown server id must error") + } + + manager.Close() + if result := manager.Call(context.Background(), echo, map[string]any{"text": "closed"}); !result.IsError { + t.Fatalf("calls after Close must fail: %+v", result) + } +} + +func TestMarketAddStaysOffUntilEnabledThenConnects(t *testing.T) { + manager := stdioManager(t, map[string]ServerConfig{}) + manager.Start() + server, err := manager.AddFromMarket(MarketEntry{Name: "helper", Config: ServerConfig{Command: os.Args[0], Env: map[string]string{stdioHelperEnv: "1"}}}) + if err != nil { + t.Fatal(err) + } + if server.Scope != ScopeUser { + t.Fatalf("marketplace adds must go to user scope, got %q", server.Scope) + } + view := viewByName(t, manager, "helper") + if view.Status != StatusOff { + t.Fatalf("marketplace add must start off, got %+v", view) + } + if err := manager.Reconnect(view.Server.ID()); err == nil || !strings.Contains(err.Error(), "is off") { + t.Fatalf("reconnecting an off server = %v", err) + } + if _, err := manager.AddFromMarket(MarketEntry{Name: "helper", Config: ServerConfig{Command: "npx"}}); err == nil { + t.Fatal("adding a duplicate name must fail") + } + if err := manager.Enable(view.Server.ID()); err != nil { + t.Fatal(err) + } + waitConnected(t, manager) + if view := viewByName(t, manager, "helper"); view.Status != StatusConnected || len(view.Tools) != 2 { + t.Fatalf("enabled market server = %+v", view) + } + if err := manager.Enable("unknown"); err == nil { + t.Fatal("unknown server id must error") + } + if err := manager.Login(context.Background(), "unknown", LoginOptions{}); err == nil { + t.Fatal("login for an unknown server must error") + } +} + +func TestProbeStatusReportsUnauthenticatedStatus(t *testing.T) { + server := httptest.NewServer(http.HandlerFunc(func(writer http.ResponseWriter, _ *http.Request) { + writer.WriteHeader(http.StatusUnauthorized) + })) + defer server.Close() + status, err := ProbeStatus(context.Background(), nil, ServerConfig{Type: TransportHTTP, URL: server.URL}) + if err != nil || status != http.StatusUnauthorized { + t.Fatalf("ProbeStatus = %d, %v", status, err) + } + if _, err := ProbeStatus(context.Background(), server.Client(), ServerConfig{Type: TransportHTTP, URL: "http://127.0.0.1:1"}); err == nil { + t.Fatal("unreachable server must error") + } +} diff --git a/internal/mcp/manager.go b/internal/mcp/manager.go new file mode 100644 index 0000000..0f0fefd --- /dev/null +++ b/internal/mcp/manager.go @@ -0,0 +1,1151 @@ +package mcp + +import ( + "bytes" + "context" + "encoding/json" + "errors" + "fmt" + "io" + "net/http" + "net/url" + "os" + "os/exec" + "path/filepath" + "sort" + "strings" + "sync" + "time" + "unicode/utf8" + + sdk "github.com/modelcontextprotocol/go-sdk/mcp" + + "github.com/rootlyhq/rootly-cli/internal/mcp/builtin" +) + +type Status string + +const ( + StatusOff Status = "off" + StatusPending Status = "pending" + StatusConnecting Status = "connecting" + StatusConnected Status = "connected" + StatusNeedsAuth Status = "auth" + StatusError Status = "error" +) + +const builtinServerID = "builtin" + +const ( + MaxClientTools = 64 + MaxSchemaBytes = 16 * 1024 + MaxClientToolsBytes = 256 * 1024 + MaxToolOutputBytes = 100 * 1024 + maxDescriptionLength = 1024 + connectTimeout = 45 * time.Second +) + +type ToolInfo struct { + Name string + Description string + ReadOnly bool + Enabled bool +} + +type BuiltinToolInfo struct { + Name string + Description string + ReadOnly bool + Enabled bool +} + +type ServerView struct { + Server Server + Status Status + Detail string + Tools []ToolInfo + Prompts []PromptInfo + Resources []ResourceInfo + SessionOff bool +} + +type ClientTool struct { + Name string `json:"name"` + Description string `json:"description"` + InputSchema json.RawMessage `json:"input_schema"` + ReadOnly bool `json:"read_only"` +} + +type ToolRef struct { + ServerID string + ServerName string + Tool string + ReadOnly bool +} + +type CallResult struct { + Output string + IsError bool +} + +type runtimeServer struct { + server Server + status Status + detail string + tools []*sdk.Tool + prompts []PromptInfo + resources []ResourceInfo + session *sdk.ClientSession + sessionOff bool + generation int + cancel context.CancelFunc +} + +type ManagerOptions struct { + CWD string + Store *Store + State *StateStore + Tokens *TokenStore + Discovery *Discovery + Version string + Lookup func(string) (string, bool) + LookPath builtin.LookPathFunc +} + +// Manager owns local MCP connections for one agent run. +type Manager struct { + options ManagerOptions + ctx context.Context + cancel context.CancelFunc + + mu sync.Mutex + servers []*runtimeServer + byID map[string]*runtimeServer + toolRefs map[string]ToolRef + builtins *builtin.Registry + builtinEnabled map[string]bool + builtinErr error + sessionTools map[string]bool + problems []string + changes chan struct{} + closed bool +} + +func NewManager(options ManagerOptions) *Manager { + if options.Store == nil { + options.Store = NewStore(options.CWD) + } + if options.State == nil { + options.State = NewStateStore() + } + if options.Tokens == nil { + options.Tokens = NewTokenStore() + } + if options.Discovery == nil { + discovery := NewDiscovery(options.CWD) + options.Discovery = &discovery + } + if options.Lookup == nil { + options.Lookup = os.LookupEnv + } + rulesDir := builtin.DefaultRulesDir() + if options.Discovery.Home != "" { + rulesDir = filepath.Join(options.Discovery.Home, ".rootly-cli", "rules") + } + builtins, builtinErr := builtin.NewRegistryWithPolicy(options.CWD, options.LookPath, builtin.LoadPolicy(rulesDir)) + ctx, cancel := context.WithCancel(context.Background()) + return &Manager{ + options: options, + ctx: ctx, + cancel: cancel, + byID: map[string]*runtimeServer{}, + toolRefs: map[string]ToolRef{}, + builtins: builtins, + builtinEnabled: map[string]bool{}, + builtinErr: builtinErr, + sessionTools: map[string]bool{}, + changes: make(chan struct{}, 1), + } +} + +// Changes signals (coalesced) whenever a server's status or tools change. +func (m *Manager) Changes() <-chan struct{} { + return m.changes +} + +func (m *Manager) notify() { + select { + case m.changes <- struct{}{}: + default: + } +} + +// Load reads every config source; servers shadowed by a higher-precedence +// scope (local > project > user) are dropped like Claude Code does. +func (m *Manager) Load() error { + configured, err := m.options.Store.All() + imported, importProblems := m.options.Discovery.Imported() + state, stateErr := m.options.State.Load() + m.mu.Lock() + defer m.mu.Unlock() + m.problems = nil + if err != nil { + m.problems = append(m.problems, err.Error()) + } + if m.builtinErr != nil { + m.problems = append(m.problems, m.builtinErr.Error()) + } + if m.builtins != nil { + if err := m.builtins.Policy().LoadError(); err != nil { + m.problems = append(m.problems, "Local tool rules (~/.rootly-cli/rules) skipped: "+err.Error()) + } + } + if stateErr != nil { + m.problems = append(m.problems, stateErr.Error()) + } + m.builtinEnabled = make(map[string]bool, len(state.BuiltinTools)) + for name, enabled := range state.BuiltinTools { + m.builtinEnabled[name] = enabled + } + for _, problem := range importProblems { + m.problems = append(m.problems, fmt.Sprintf("%s (%s): %v", problem.Source.Label, problem.Source.Path, problem.Err)) + } + seen := map[string]bool{} + var servers []Server + for _, server := range configured { + if seen[server.Name] { + continue + } + seen[server.Name] = true + servers = append(servers, server) + } + servers = append(servers, imported...) + previous := m.byID + m.servers = nil + m.byID = map[string]*runtimeServer{} + for _, server := range servers { + runtime := previous[server.ID()] + if runtime == nil || runtime.server.Config.Hash() != server.Config.Hash() || runtime.server.Project != server.Project { + if runtime != nil { + runtime.close() + } + runtime = &runtimeServer{server: server} + runtime.status = initialStatus(state, server) + } else if status := initialStatus(state, server); status == StatusOff || status == StatusPending { + runtime.close() + runtime.status, runtime.detail = status, "" + } + m.servers = append(m.servers, runtime) + m.byID[server.ID()] = runtime + } + for id, runtime := range previous { + if m.byID[id] == nil { + runtime.close() + } + } + m.notify() + if len(m.problems) > 0 { + return errors.New(strings.Join(m.problems, "; ")) + } + return nil +} + +func initialStatus(state State, server Server) Status { + switch { + case !state.IsEnabled(server): + return StatusOff + case state.NeedsApproval(server): + return StatusPending + default: + return StatusConnecting + } +} + +func (m *Manager) Problems() []string { + m.mu.Lock() + defer m.mu.Unlock() + return append([]string(nil), m.problems...) +} + +// Start connects every enabled, approved server in the background. +func (m *Manager) Start() { + m.mu.Lock() + defer m.mu.Unlock() + for _, runtime := range m.servers { + if runtime.status == StatusConnecting && runtime.session == nil && !runtime.sessionOff { + m.connectLocked(runtime) + } + } +} + +func (m *Manager) connectLocked(runtime *runtimeServer) { + if m.closed { + return + } + runtime.close() + runtime.generation++ + generation := runtime.generation + runtime.status = StatusConnecting + runtime.detail = "" + ctx, cancel := context.WithCancel(m.ctx) + runtime.cancel = cancel + server := runtime.server + go func() { + session, tools, status, detail, probe := m.connect(ctx, server) + var prompts []PromptInfo + var resources []ResourceInfo + if session != nil { + discoverCtx, cancelDiscover := context.WithTimeout(ctx, server.Config.CallTimeout(connectTimeout)) + prompts, resources = discoverExtras(discoverCtx, session) + cancelDiscover() + } + m.mu.Lock() + defer m.mu.Unlock() + if runtime.generation != generation || m.closed { + if session != nil { + _ = session.Close() + } + return + } + runtime.session, runtime.tools, runtime.status, runtime.detail = session, tools, status, detail + runtime.prompts, runtime.resources = prompts, resources + if session != nil { + m.watchAuth(runtime, generation, probe) + } + m.notify() + }() + m.notify() +} + +// watchAuth must run with m.mu held; a 401 seen while discovering prompts and resources marks the server at once. +func (m *Manager) watchAuth(runtime *runtimeServer, generation int, probe *authProbe) { + if probe.watch(func() { m.loseAuth(runtime, generation) }) { + runtime.close() + runtime.status, runtime.detail = StatusNeedsAuth, "needs authentication, Enter to log in" + } +} + +func (m *Manager) loseAuth(runtime *runtimeServer, generation int) { + m.mu.Lock() + defer m.mu.Unlock() + if runtime.generation != generation || runtime.session == nil { + return + } + runtime.close() + runtime.status, runtime.detail = StatusNeedsAuth, "needs authentication, Enter to log in" + m.notify() +} + +func (r *runtimeServer) close() { + r.generation++ + if r.cancel != nil { + r.cancel() + r.cancel = nil + } + if r.session != nil { + session := r.session + r.session = nil + go func() { _ = session.Close() }() + } + r.tools = nil + r.prompts = nil + r.resources = nil +} + +func (m *Manager) connect(ctx context.Context, server Server) (session *sdk.ClientSession, tools []*sdk.Tool, status Status, detail string, probe *authProbe) { + config := server.Config.Expanded(m.options.Lookup) + connectCtx, cancel := context.WithTimeout(ctx, config.CallTimeout(connectTimeout)) + defer cancel() + transport, stderr, probe, err := m.transport(ctx, server, config) + if err != nil { + return nil, nil, StatusError, err.Error(), nil + } + client := sdk.NewClient(&sdk.Implementation{Name: "rootly-cli", Version: m.options.Version}, nil) + session, err = client.Connect(connectCtx, transport, nil) + if err != nil { + return nil, nil, failureStatus(err, probe), failureDetail(err, stderr, probe), nil + } + for tool, err := range session.Tools(connectCtx, nil) { + if err != nil { + _ = session.Close() + return nil, nil, failureStatus(err, probe), failureDetail(err, stderr, probe), nil + } + tools = append(tools, tool) + } + sort.SliceStable(tools, func(i, j int) bool { return tools[i].Name < tools[j].Name }) + return session, tools, StatusConnected, "", probe +} +func failureStatus(err error, unauthorized *authProbe) Status { + if errors.Is(err, ErrNeedsAuth) || unauthorized.seen() { + return StatusNeedsAuth + } + return StatusError +} + +func failureDetail(err error, stderr *tailBuffer, unauthorized *authProbe) string { + if errors.Is(err, ErrNeedsAuth) || unauthorized.seen() { + return "needs authentication, Enter to log in" + } + if errors.Is(err, exec.ErrNotFound) { + return "exited: command not found" + } + message := err.Error() + if stderr != nil { + if line := stderr.lastLine(); line != "" { + message = line + } + } + return oneLine(message) +} + +func oneLine(value string) string { + value = strings.Join(strings.Fields(value), " ") + if len(value) > 160 { + value = value[:157] + "..." + } + return value +} + +func (m *Manager) transport(ctx context.Context, server Server, config ServerConfig) (sdk.Transport, *tailBuffer, *authProbe, error) { + switch config.Transport() { + case TransportStdio: + if _, err := exec.LookPath(config.Command); err != nil { + return nil, nil, nil, fmt.Errorf("exited: command not found") + } + command := exec.CommandContext(ctx, config.Command, config.Args...) + command.Env = os.Environ() + for _, key := range sortedKeys(config.Env) { + command.Env = append(command.Env, key+"="+config.Env[key]) + } + if server.Scope == ScopeProject || server.Scope == ScopeLocal { + command.Dir = m.options.CWD + } + stderr := &tailBuffer{} + command.Stderr = stderr + return &sdk.CommandTransport{Command: command, TerminateDuration: 2 * time.Second}, stderr, nil, nil + case TransportHTTP, TransportSSE: + probe := &authProbe{} + bearer, refresh := m.bearerFuncs(config) + client := newServerHTTPClientWithRefresh(config, bearer, refresh, probe) + if config.Transport() == TransportSSE { + return &sdk.SSEClientTransport{Endpoint: config.URL, HTTPClient: client}, nil, probe, nil + } + return &sdk.StreamableClientTransport{Endpoint: config.URL, HTTPClient: client, DisableStandaloneSSE: true}, nil, probe, nil + default: + return nil, nil, nil, fmt.Errorf("unsupported transport %q", config.Type) + } +} + +func (m *Manager) bearerFuncs(config ServerConfig) (bearer, refresh func(context.Context) string) { + if config.BearerTokenEnvVar != "" { + return func(context.Context) string { + value, _ := m.options.Lookup(config.BearerTokenEnvVar) + return value + }, nil + } + tokens := m.options.Tokens + serverURL := config.URL + var source *persistingTokenSource + var mu sync.Mutex + bearer = func(ctx context.Context) string { + mu.Lock() + defer mu.Unlock() + if source == nil || !source.base.token().Valid() { + credential, ok := tokens.Get(serverURL) + if !ok { + return "" + } + source = tokens.tokenSourceWithContext(ctx, serverURL, credential, credential.oauthConfig(), credential.token(), nil).(*persistingTokenSource) + } + token, err := source.Token() + if err != nil { + return "" + } + return token.AccessToken + } + refresh = func(ctx context.Context) string { + mu.Lock() + defer mu.Unlock() + credential, ok := tokens.Get(serverURL) + if !ok || credential.RefreshToken == "" { + return "" + } + expired := credential.token() + expired.Expiry = time.Now().Add(-time.Minute) + source = nil + refreshing := &persistingTokenSource{store: tokens, serverURL: serverURL, base: credential, source: credential.oauthConfig().TokenSource(ctx, expired), last: expired.AccessToken} + token, err := refreshing.Token() + if err != nil || token.AccessToken == "" || token.AccessToken == credential.AccessToken { + return "" + } + return token.AccessToken + } + return bearer, refresh +} + +type authProbe struct { + mu sync.Mutex + unauthorized bool + onExpired func() +} + +func (p *authProbe) mark() { + p.mu.Lock() + p.unauthorized = true + onExpired := p.onExpired + p.mu.Unlock() + if onExpired != nil { + onExpired() + } +} + +func (p *authProbe) watch(onExpired func()) (alreadyUnauthorized bool) { + if p == nil { + return false + } + p.mu.Lock() + defer p.mu.Unlock() + p.onExpired = onExpired + return p.unauthorized +} + +func (p *authProbe) seen() bool { + if p == nil { + return false + } + p.mu.Lock() + defer p.mu.Unlock() + return p.unauthorized +} + +// newServerHTTPClient sends the configured headers and bearer token only to +// the server's own origin and refuses redirects to any other origin. +func newServerHTTPClient(config ServerConfig, bearer func(context.Context) string, probe *authProbe) *http.Client { + return newServerHTTPClientWithRefresh(config, bearer, nil, probe) +} + +func newServerHTTPClientWithRefresh(config ServerConfig, bearer, refresh func(context.Context) string, probe *authProbe) *http.Client { + return &http.Client{ + Transport: &headerTransport{ + base: http.DefaultTransport, + origin: originOf(config.URL), + headers: config.Headers, + bearer: bearer, + refresh: refresh, + probe: probe, + }, + CheckRedirect: rejectCrossOriginRedirect, + } +} + +func originOf(raw string) string { + parsed, err := url.Parse(raw) + if err != nil || parsed.Host == "" { + return "" + } + scheme, port := strings.ToLower(parsed.Scheme), parsed.Port() + if port == "" { + port = map[string]string{"http": "80", "https": "443"}[scheme] + } + return scheme + "://" + strings.ToLower(parsed.Hostname()) + ":" + port +} + +func rejectCrossOriginRedirect(request *http.Request, via []*http.Request) error { + if len(via) >= 10 { + return errors.New("stopped after 10 redirects") + } + if len(via) > 0 && originOf(request.URL.String()) != originOf(via[0].URL.String()) { + return fmt.Errorf("refusing cross-origin redirect to %s", request.URL.Host) + } + return nil +} + +type headerTransport struct { + base http.RoundTripper + origin string + headers map[string]string + bearer func(context.Context) string + refresh func(context.Context) string + probe *authProbe +} + +func (t *headerTransport) RoundTrip(request *http.Request) (*http.Response, error) { + if t.origin == "" || originOf(request.URL.String()) != t.origin { + return t.base.RoundTrip(request) + } + request = request.Clone(request.Context()) + for key, value := range t.headers { + request.Header.Set(key, value) + } + usedBearer := false + if request.Header.Get("Authorization") == "" && t.bearer != nil { + if token := t.bearer(request.Context()); token != "" { + request.Header.Set("Authorization", "Bearer "+token) + usedBearer = true + } + } + response, err := t.base.RoundTrip(request) + if err == nil && response.StatusCode == http.StatusUnauthorized && usedBearer { + response, err = t.retryWithRefreshedToken(request, response) + } + if err == nil && response.StatusCode == http.StatusUnauthorized && t.probe != nil { + t.probe.mark() + } + return response, err +} + +type tailBuffer struct { + mu sync.Mutex + data []byte +} + +func (b *tailBuffer) Write(p []byte) (int, error) { + b.mu.Lock() + defer b.mu.Unlock() + b.data = append(b.data, p...) + if len(b.data) > 4096 { + b.data = b.data[len(b.data)-4096:] + } + return len(p), nil +} + +func (b *tailBuffer) lastLine() string { + b.mu.Lock() + defer b.mu.Unlock() + lines := bytes.Split(bytes.TrimSpace(b.data), []byte("\n")) + if len(lines) == 0 { + return "" + } + return strings.TrimSpace(string(lines[len(lines)-1])) +} + +func (m *Manager) Views() []ServerView { + m.mu.Lock() + defer m.mu.Unlock() + views := make([]ServerView, 0, len(m.servers)) + for _, runtime := range m.servers { + view := ServerView{Server: runtime.server, Status: runtime.status, Detail: runtime.detail, SessionOff: runtime.sessionOff, Prompts: runtime.prompts, Resources: runtime.resources} + view.Tools = make([]ToolInfo, 0, len(runtime.tools)) + for _, tool := range runtime.tools { + view.Tools = append(view.Tools, ToolInfo{Name: tool.Name, Description: tool.Description, ReadOnly: readOnly(tool)}) + } + views = append(views, view) + } + return views +} + +// Connecting returns how many servers are still starting. +func (m *Manager) Connecting() int { + m.mu.Lock() + defer m.mu.Unlock() + count := 0 + for _, runtime := range m.servers { + if runtime.status == StatusConnecting && !runtime.sessionOff { + count++ + } + } + return count +} + +func (m *Manager) withServer(id string, change func(*runtimeServer) error) error { + m.mu.Lock() + defer m.mu.Unlock() + runtime := m.byID[id] + if runtime == nil { + return fmt.Errorf("unknown MCP server") + } + err := change(runtime) + m.notify() + return err +} + +// ToggleSession turns a server on or off for this agent run only. +func (m *Manager) ToggleSession(id string) error { + return m.withServer(id, func(runtime *runtimeServer) error { + runtime.sessionOff = !runtime.sessionOff + if !runtime.sessionOff && runtime.session == nil && (runtime.status == StatusConnecting || runtime.status == StatusError) { + m.connectLocked(runtime) + } + return nil + }) +} + +// Approve records approval of a project server and starts it. +func (m *Manager) Approve(id string) error { + return m.withServer(id, func(runtime *runtimeServer) error { + if !runtime.server.NeedsProjectTrust() { + return errors.New("only project-controlled servers need approval") + } + if err := m.options.State.Update(func(state *State) { state.Approve(runtime.server) }); err != nil { + return err + } + if runtime.status == StatusPending { + m.connectLocked(runtime) + } + return nil + }) +} + +// Enable persistently enables a server (imported servers start off). +func (m *Manager) Enable(id string) error { + return m.withServer(id, func(runtime *runtimeServer) error { + if err := m.options.State.Update(func(state *State) { state.SetEnabled(runtime.server, true) }); err != nil { + return err + } + state, err := m.options.State.Load() + if err != nil { + return err + } + if state.NeedsApproval(runtime.server) { + runtime.status = StatusPending + return nil + } + m.connectLocked(runtime) + return nil + }) +} + +func (m *Manager) Reconnect(id string) error { + return m.withServer(id, func(runtime *runtimeServer) error { + if runtime.status == StatusOff || runtime.status == StatusPending { + return fmt.Errorf("%s is %s", runtime.server.Name, map[Status]string{StatusOff: "off", StatusPending: "pending approval"}[runtime.status]) + } + runtime.sessionOff = false + m.connectLocked(runtime) + return nil + }) +} + +// Login runs OAuth for an HTTP server, then reconnects it. +func (m *Manager) Login(ctx context.Context, id string, options LoginOptions) error { + m.mu.Lock() + runtime := m.byID[id] + m.mu.Unlock() + if runtime == nil { + return fmt.Errorf("unknown MCP server") + } + config := runtime.server.Config.Expanded(m.options.Lookup) + if err := m.options.Tokens.Login(ctx, config, options); err != nil { + return err + } + m.mu.Lock() + inactive := runtime.status == StatusOff || runtime.status == StatusPending + m.mu.Unlock() + if inactive { + m.notify() + return nil + } + return m.Reconnect(id) +} + +// ClientTools returns the tools to advertise for the next turn and +// remembers how each generated name maps back to its server and tool. +func (m *Manager) ClientTools() (advertised []ClientTool, omitted int) { + m.mu.Lock() + defer m.mu.Unlock() + namer := NewClientToolNamer() + refs := map[string]ToolRef{} + var tools []ClientTool + skipped := 0 + total := 2 + if m.builtins != nil { + builtinTools := m.builtins.Tools() + for _, tool := range builtinTools { + namer.Reserve(tool.Name()) + } + for _, tool := range builtinTools { + if !m.builtinEnabledLocked(tool.Name()) { + continue + } + schema := tool.Schema() + if len(schema) > MaxSchemaBytes || len(tools) >= MaxClientTools { + skipped++ + continue + } + candidate := ClientTool{ + Name: tool.Name(), + Description: truncateRunes(tool.Description(), maxDescriptionLength), + InputSchema: schema, + ReadOnly: tool.AdvertisedReadOnly(), + } + encoded, err := json.Marshal(candidate) + if err != nil || total+len(encoded)+1 > MaxClientToolsBytes { + skipped++ + continue + } + total += len(encoded) + 1 + tools = append(tools, candidate) + refs[candidate.Name] = ToolRef{ + ServerID: builtinServerID, + ServerName: "Built-in", + Tool: strings.TrimPrefix(candidate.Name, ClientToolPrefix), + ReadOnly: candidate.ReadOnly, + } + } + } + for _, runtime := range m.servers { + if runtime.status != StatusConnected || runtime.sessionOff { + continue + } + for _, tool := range runtime.tools { + schema, ok := clientSchema(tool.InputSchema) + if !ok || len(tools) >= MaxClientTools { + skipped++ + continue + } + candidate := ClientTool{ + Name: namer.Name(runtime.server.Name, tool.Name), + Description: truncateRunes(firstNonEmpty(tool.Description, tool.Title, tool.Name), maxDescriptionLength), + InputSchema: schema, + ReadOnly: readOnly(tool), + } + encoded, _ := json.Marshal(candidate) + if total+len(encoded)+1 > MaxClientToolsBytes { + skipped++ + continue + } + total += len(encoded) + 1 + tools = append(tools, candidate) + refs[candidate.Name] = ToolRef{ServerID: runtime.server.ID(), ServerName: runtime.server.Name, Tool: tool.Name, ReadOnly: candidate.ReadOnly} + } + } + m.toolRefs = refs + return tools, skipped +} + +func firstNonEmpty(values ...string) string { + for _, value := range values { + if strings.TrimSpace(value) != "" { + return value + } + } + return "" +} + +func truncateRunes(value string, limit int) string { + if utf8.RuneCountInString(value) <= limit { + return value + } + runes := []rune(value) + return string(runes[:limit-1]) + "…" +} + +func readOnly(tool *sdk.Tool) bool { + return tool.Annotations != nil && tool.Annotations.ReadOnlyHint +} + +func clientSchema(schema any) (json.RawMessage, bool) { + if schema == nil { + return json.RawMessage(`{"type":"object","properties":{}}`), true + } + encoded, err := json.Marshal(schema) + if err != nil || len(encoded) > MaxSchemaBytes { + return nil, false + } + var object map[string]any + if json.Unmarshal(encoded, &object) != nil || object == nil { + return nil, false + } + if kind, ok := object["type"]; ok && kind != "object" { + return nil, false + } + if _, ok := object["type"]; !ok { + object["type"] = "object" + encoded, _ = json.Marshal(object) + } + return encoded, true +} + +func (m *Manager) Lookup(name string) (ToolRef, bool) { + m.mu.Lock() + defer m.mu.Unlock() + ref, ok := m.toolRefs[name] + return ref, ok +} + +func (m *Manager) NeedsApproval(name string, args map[string]any) bool { + m.mu.Lock() + defer m.mu.Unlock() + ref, ok := m.toolRefs[name] + if !ok { + return true + } + if ref.ServerID != builtinServerID { + return !ref.ReadOnly && !m.sessionTools[sessionToolKey(ref)] + } + tool, ok := m.builtins.Lookup(name) + return !ok || !m.builtinEnabledLocked(name) || tool.NeedsApproval(args) +} + +// sessionToolKey identifies an MCP tool by its server and original name: +// generated client names are rebuilt each turn and can move between servers. +func sessionToolKey(ref ToolRef) string { + return ref.ServerID + "\x00" + ref.Tool +} + +// ReadOnlyCall reports whether a call reads only on its own merits (not +// through a remembered approval or allow rule), so it may run alongside +// other reads. +func (m *Manager) ReadOnlyCall(name string, args map[string]any) bool { + m.mu.Lock() + defer m.mu.Unlock() + ref, ok := m.toolRefs[name] + if !ok { + return false + } + if ref.ServerID != builtinServerID { + return ref.ReadOnly + } + tool, ok := m.builtins.Lookup(name) + return ok && m.builtinEnabledLocked(name) && tool.Approval(args).ReadOnly +} + +// ApprovalOptions are the ways a call that needs approval may be remembered: +// for the session (MCP tools by name, commands by Prefix) or, for commands, +// as an allow rule in ~/.rootly-cli/rules/default.rules. +type ApprovalOptions struct { + Session bool + Persist bool + Prefix []string +} + +func (m *Manager) ApprovalOptions(name string, args map[string]any) ApprovalOptions { + m.mu.Lock() + defer m.mu.Unlock() + ref, ok := m.toolRefs[name] + if !ok { + return ApprovalOptions{} + } + if ref.ServerID != builtinServerID { + return ApprovalOptions{Session: !ref.ReadOnly} + } + tool, ok := m.builtins.Lookup(name) + if !ok || !m.builtinEnabledLocked(name) { + return ApprovalOptions{} + } + approval := tool.Approval(args) + if !approval.Needs || len(approval.Prefix) == 0 { + return ApprovalOptions{} + } + return ApprovalOptions{Session: true, Persist: true, Prefix: approval.Prefix} +} + +// RememberApproval allows this kind of call for the session or, with +// persist, saves an allow rule for its command prefix. +func (m *Manager) RememberApproval(name string, args map[string]any, persist bool) error { + options := m.ApprovalOptions(name, args) + switch { + case len(options.Prefix) > 0 && persist: + return m.builtins.Policy().SaveAllow(options.Prefix) + case len(options.Prefix) > 0: + m.builtins.Policy().AllowForSession(options.Prefix) + return nil + case options.Session && !persist: + m.mu.Lock() + defer m.mu.Unlock() + ref, ok := m.toolRefs[name] + if !ok { + return errors.New("unknown local tool " + name) + } + m.sessionTools[sessionToolKey(ref)] = true + return nil + } + return errors.New("this call can only be approved once") +} + +func (m *Manager) BuiltinTools() []BuiltinToolInfo { + m.mu.Lock() + defer m.mu.Unlock() + if m.builtins == nil { + return nil + } + tools := m.builtins.Tools() + result := make([]BuiltinToolInfo, 0, len(tools)) + for _, tool := range tools { + result = append(result, BuiltinToolInfo{ + Name: tool.Name(), + Description: tool.Description(), + ReadOnly: tool.AdvertisedReadOnly(), + Enabled: m.builtinEnabledLocked(tool.Name()), + }) + } + return result +} + +func (m *Manager) SetBuiltinEnabled(name string, enabled bool) error { + if m.builtins == nil { + return errors.New("built-in tools are unavailable") + } + if _, ok := m.builtins.Lookup(name); !ok { + return fmt.Errorf("unknown built-in tool %q", name) + } + if err := m.options.State.Update(func(state *State) { state.SetBuiltinEnabled(name, enabled) }); err != nil { + return err + } + m.mu.Lock() + m.builtinEnabled[name] = enabled + m.mu.Unlock() + m.notify() + return nil +} + +func (m *Manager) builtinEnabledLocked(name string) bool { + if enabled, ok := m.builtinEnabled[name]; ok { + return enabled + } + return name != "local_bash" +} + +// Call runs a client tool by its generated name. +func (m *Manager) Call(ctx context.Context, name string, arguments map[string]any) CallResult { + m.mu.Lock() + ref, ok := m.toolRefs[name] + if ok && ref.ServerID == builtinServerID { + tool, found := m.builtins.Lookup(name) + enabled := m.builtinEnabledLocked(name) + m.mu.Unlock() + if !found { + return CallResult{Output: "Unknown built-in tool " + name, IsError: true} + } + if !enabled { + return CallResult{Output: "Built-in tool " + name + " is off", IsError: true} + } + result := tool.Run(ctx, arguments) + return CallResult{Output: TruncateOutput(result.Output), IsError: result.IsError} + } + var session *sdk.ClientSession + timeout := time.Duration(0) + if ok { + if runtime := m.byID[ref.ServerID]; runtime != nil && !runtime.sessionOff { + session = runtime.session + timeout = runtime.server.Config.CallTimeout(0) + } + } + m.mu.Unlock() + if !ok { + return CallResult{Output: "Unknown local tool " + name, IsError: true} + } + if session == nil { + return CallResult{Output: "Local MCP server " + ref.ServerName + " is not connected", IsError: true} + } + if arguments == nil { + arguments = map[string]any{} + } + if timeout > 0 { + var cancel context.CancelFunc + ctx, cancel = context.WithTimeout(ctx, timeout) + defer cancel() + } + result, err := session.CallTool(ctx, &sdk.CallToolParams{Name: ref.Tool, Arguments: arguments}) + if err != nil && timeout > 0 && errors.Is(ctx.Err(), context.DeadlineExceeded) { + return CallResult{Output: fmt.Sprintf("Local tool timed out after %s", timeout), IsError: true} + } + if err != nil { + return CallResult{Output: TruncateOutput("Local tool failed: " + err.Error()), IsError: true} + } + return CallResult{Output: TruncateOutput(renderResult(result)), IsError: result.IsError} +} + +func renderResult(result *sdk.CallToolResult) string { + var parts []string + for _, content := range result.Content { + if text, ok := content.(*sdk.TextContent); ok { + parts = append(parts, text.Text) + continue + } + encoded, err := json.Marshal(content) + if err == nil { + parts = append(parts, string(encoded)) + } + } + if len(parts) == 0 && result.StructuredContent != nil { + if encoded, err := json.Marshal(result.StructuredContent); err == nil { + parts = append(parts, string(encoded)) + } + } + return strings.Join(parts, "\n") +} + +const truncationMarker = "\n[output truncated to 100 KB]" + +// TruncateOutput caps output at 100 KB on a UTF-8 boundary. +func TruncateOutput(output string) string { + if len(output) <= MaxToolOutputBytes { + return output + } + cut := MaxToolOutputBytes - len(truncationMarker) + for cut > 0 && !utf8.RuneStart(output[cut]) { + cut-- + } + return output[:cut] + truncationMarker +} + +// Close stops every server (stdio processes are terminated). +func (m *Manager) Close() { + m.mu.Lock() + m.closed = true + sessions := make([]*sdk.ClientSession, 0, len(m.servers)) + for _, runtime := range m.servers { + if runtime.session != nil { + sessions = append(sessions, runtime.session) + runtime.session = nil + } + } + m.mu.Unlock() + var wg sync.WaitGroup + for _, session := range sessions { + wg.Add(1) + go func() { + defer wg.Done() + _ = session.Close() + }() + } + done := make(chan struct{}) + go func() { + wg.Wait() + close(done) + }() + select { + case <-done: + case <-time.After(3 * time.Second): + } + m.cancel() +} + +func (t *headerTransport) retryWithRefreshedToken(request *http.Request, response *http.Response) (*http.Response, error) { + if t.refresh == nil || (request.Body != nil && request.Body != http.NoBody && request.GetBody == nil) { + return response, nil + } + token := t.refresh(request.Context()) + if token == "" { + return response, nil + } + retry := request.Clone(request.Context()) + if request.GetBody != nil { + body, err := request.GetBody() + if err != nil { + return response, nil + } + retry.Body = body + } + retry.Header.Set("Authorization", "Bearer "+token) + _, _ = io.Copy(io.Discard, response.Body) + _ = response.Body.Close() + return t.base.RoundTrip(retry) +} + +// AddFromMarket installs a marketplace entry in USER scope, switched off until the user turns it on. +func (m *Manager) AddFromMarket(entry MarketEntry) (Server, error) { + imported, _ := m.options.Discovery.Imported() + for _, server := range imported { + if server.Name == entry.Name { + return Server{}, fmt.Errorf("an MCP server named %s is already imported from %s", entry.Name, server.Path) + } + } + server, err := AddFromMarket(m.options.Store, m.options.State, entry) + if err != nil { + return Server{}, err + } + _ = m.Load() + return server, nil +} diff --git a/internal/mcp/mask.go b/internal/mcp/mask.go new file mode 100644 index 0000000..d1649e9 --- /dev/null +++ b/internal/mcp/mask.go @@ -0,0 +1,220 @@ +package mcp + +import ( + "net/url" + "regexp" + "sort" + "strings" +) + +const Mask = "*****" + +// MaskedConfig hides env/header values and URL credentials for display. +func (c ServerConfig) Masked() ServerConfig { + masked := c + masked.Args = maskArgs(c.Args, false) + masked.URL = maskURL(c.URL) + if len(c.Env) > 0 { + masked.Env = make(map[string]string, len(c.Env)) + for key := range c.Env { + masked.Env[key] = Mask + } + } + if len(c.Headers) > 0 { + masked.Headers = make(map[string]string, len(c.Headers)) + for key := range c.Headers { + masked.Headers[key] = Mask + } + } + return masked +} + +func maskURL(raw string) string { + parsed, err := url.Parse(raw) + if err != nil || parsed.Host == "" { + return raw + } + if parsed.User != nil { + parsed.User = url.User(Mask) + } + if parsed.RawQuery != "" { + query := parsed.Query() + for key := range query { + query.Set(key, Mask) + } + parsed.RawQuery = query.Encode() + } + return strings.ReplaceAll(parsed.String(), url.QueryEscape(Mask), Mask) +} + +// Origin is a one-line, secret-free description of where a server runs. +func (c ServerConfig) Origin() string { + if c.Transport() == TransportStdio { + return strings.TrimSpace(c.Command + " " + strings.Join(maskArgs(c.Args, false), " ")) + } + return maskURL(c.URL) +} + +func sortedKeys(values map[string]string) []string { + keys := make([]string, 0, len(values)) + for key := range values { + keys = append(keys, key) + } + sort.Strings(keys) + return keys +} + +// MaskedPairs renders KEY=***** pairs in key order. +func MaskedPairs(values map[string]string, separator string) []string { + pairs := make([]string, 0, len(values)) + for _, key := range sortedKeys(values) { + pairs = append(pairs, key+separator+Mask) + } + return pairs +} + +var secretArgName = regexp.MustCompile(`(?i)(token|secret|password|passwd|api[-_]?key|apikey|auth|credential|bearer|access[-_]?key|private[-_]?key)`) +var secretCommandArgument = regexp.MustCompile(`(?i)(--?[\w-]*(?:token|secret|password|passwd|api[-_]?key|apikey|auth|credential|access[-_]?key|private[-_]?key)[\w-]*(?:=|\s+)|\b[A-Z_]*(?:TOKEN|SECRET|PASSWORD|PASSWD|API_KEY|ACCESS_KEY|PRIVATE_KEY)[A-Z_]*=)("[^"]*"|'[^']*'|[^\s;&|]+)`) +var secretEnvironmentArgument = regexp.MustCompile(`(?i)^[A-Z_]*(?:TOKEN|SECRET|PASSWORD|PASSWD|API_KEY|ACCESS_KEY|PRIVATE_KEY)[A-Z_]*=`) +var bearerCommandValue = regexp.MustCompile(`(?i)\bBearer\s+[^\s"';&|]+`) + +// maskArgs hides values of secret-looking flags (--token=x, --api-key x) +// and inline Bearer tokens. +func maskArgs(args []string, allowShellCode bool) []string { + if len(args) == 0 { + return args + } + masked := make([]string, len(args)) + hideNext := false + for index, arg := range args { + switch { + case hideNext: + if allowShellCode && containsShellCode(arg) { + masked[index] = arg + } else { + masked[index] = Mask + } + hideNext = false + case allowShellCode && containsShellCode(arg): + masked[index] = arg + case secretEnvironmentArgument.MatchString(arg): + name, _, _ := strings.Cut(arg, "=") + masked[index] = name + "=" + Mask + case strings.HasPrefix(arg, "-") && strings.Contains(arg, "="): + name, _, _ := strings.Cut(arg, "=") + if secretArgName.MatchString(name) { + masked[index] = name + "=" + Mask + } else { + masked[index] = arg + } + case strings.HasPrefix(arg, "-") && secretArgName.MatchString(arg): + masked[index] = arg + hideNext = true + case strings.HasPrefix(strings.ToLower(arg), "bearer "): + masked[index] = arg[:len("bearer ")] + Mask + default: + masked[index] = arg + } + } + return masked +} + +func MaskCommandArgs(args []string) []string { + return maskArgs(args, false) +} + +func MaskCommand(command string) string { + return maskCommand(command, false) +} + +// MaskApprovalCommand keeps executable code visible because approval cards must show it. +func MaskApprovalCommand(command string) string { + return maskCommand(command, true) +} + +func maskCommand(command string, allowShellCode bool) string { + command = secretCommandArgument.ReplaceAllStringFunc(command, func(argument string) string { + parts := secretCommandArgument.FindStringSubmatch(argument) + if allowShellCode && containsShellCode(parts[2]) { + return argument + } + return parts[1] + Mask + }) + return bearerCommandValue.ReplaceAllStringFunc(command, func(argument string) string { + if allowShellCode && containsShellCode(argument) { + return argument + } + return "Bearer " + Mask + }) +} + +func MaskToolArguments(args map[string]any) map[string]any { + return maskToolArguments(args, false) +} + +// MaskApprovalArguments keeps executable code visible because approval cards must show it. +func MaskApprovalArguments(args map[string]any) map[string]any { + return maskToolArguments(args, true) +} + +func maskToolArguments(args map[string]any, allowShellCode bool) map[string]any { + if args == nil { + return nil + } + masked := make(map[string]any, len(args)) + for key, value := range args { + masked[key] = maskToolArgument(key, value, allowShellCode) + } + return masked +} + +func maskToolArgument(key string, value any, allowShellCode bool) any { + if text, ok := value.(string); ok && allowShellCode && containsShellCode(text) { + return text + } + if secretArgName.MatchString(key) { + return Mask + } + switch item := value.(type) { + case map[string]any: + return maskToolArguments(item, allowShellCode) + case []string: + if key == "args" { + return maskArgs(item, allowShellCode) + } + masked := make([]string, len(item)) + for index, text := range item { + masked[index] = maskCommand(text, allowShellCode) + } + return masked + case []any: + if key == "args" { + args := make([]string, len(item)) + for index, argument := range item { + if text, ok := argument.(string); ok { + args[index] = text + continue + } + args[index] = Mask + } + return maskArgs(args, allowShellCode) + } + masked := make([]any, len(item)) + for index, nested := range item { + masked[index] = maskToolArgument("", nested, allowShellCode) + } + return masked + case string: + return maskCommand(item, allowShellCode) + default: + return value + } +} + +func containsShellCode(value string) bool { + return strings.Contains(value, "$(") || + strings.Contains(value, "<(") || + strings.Contains(value, ">(") || + strings.ContainsRune(value, '`') || + strings.Contains(value, "${") +} diff --git a/internal/mcp/mcp_test.go b/internal/mcp/mcp_test.go new file mode 100644 index 0000000..b54305b --- /dev/null +++ b/internal/mcp/mcp_test.go @@ -0,0 +1,353 @@ +package mcp + +import ( + "encoding/json" + "os" + "path/filepath" + "runtime" + "strings" + "testing" +) + +func writeFile(t *testing.T, path, content string) { + t.Helper() + if err := os.MkdirAll(filepath.Dir(path), 0o755); err != nil { + t.Fatal(err) + } + if err := os.WriteFile(path, []byte(content), 0o600); err != nil { + t.Fatal(err) + } +} + +func readJSON(t *testing.T, path string) map[string]any { + t.Helper() + data, err := os.ReadFile(path) + if err != nil { + t.Fatal(err) + } + var document map[string]any + if err := json.Unmarshal(data, &document); err != nil { + t.Fatal(err) + } + return document +} + +func marshalTestJSON(t *testing.T, value any) string { + t.Helper() + data, err := json.Marshal(value) + if err != nil { + t.Fatal(err) + } + return string(data) +} + +func assertFileModePerm(t *testing.T, path string, want os.FileMode) { + t.Helper() + info, err := os.Stat(path) + if err != nil { + t.Fatal(err) + } + if runtime.GOOS != "windows" && info.Mode().Perm() != want { + t.Fatalf("file permissions = %v, want %v", info.Mode().Perm(), want) + } +} + +func TestStoreScopesWriteClaudeShapes(t *testing.T) { + home, project := t.TempDir(), t.TempDir() + store := &Store{UserPath: filepath.Join(home, ".rootly-cli", "mcp.json"), CWD: project} + writeFile(t, filepath.Join(project, ".mcp.json"), `{"mcpServers":{"maestro":{"command":"maestro","args":["mcp"]}},"other":{"keep":true}}`) + + if err := store.Add(ScopeLocal, "fs", ServerConfig{Command: "npx", Args: []string{"-y", "server"}, Env: map[string]string{"TOKEN": "secret"}}); err != nil { + t.Fatal(err) + } + if err := store.Add(ScopeUser, "braintrust", ServerConfig{Type: TransportHTTP, URL: "https://api.braintrust.dev/mcp"}); err != nil { + t.Fatal(err) + } + if err := store.Add(ScopeProject, "dart", ServerConfig{Command: "fvm", Args: []string{"dart", "mcp-server"}}); err != nil { + t.Fatal(err) + } + if err := store.Add(ScopeUser, "braintrust", ServerConfig{Type: TransportHTTP, URL: "https://x"}); err == nil { + t.Fatal("duplicate name in the same scope should fail") + } + + user := readJSON(t, store.UserPath) + if user["mcpServers"].(map[string]any)["braintrust"].(map[string]any)["type"] != "http" { + t.Fatalf("user scope: %v", user) + } + local := user["projects"].(map[string]any)[project].(map[string]any)["mcpServers"].(map[string]any)["fs"].(map[string]any) + if local["command"] != "npx" || local["type"] != nil { + t.Fatalf("local scope should be stdio without type: %v", local) + } + projectDocument := readJSON(t, filepath.Join(project, ".mcp.json")) + if projectDocument["other"] == nil || projectDocument["mcpServers"].(map[string]any)["maestro"] == nil || projectDocument["mcpServers"].(map[string]any)["dart"] == nil { + t.Fatalf("project write must preserve other keys and servers: %v", projectDocument) + } + assertFileModePerm(t, store.UserPath, 0o600) + + all, err := store.All() + if err != nil || len(all) != 4 { + t.Fatalf("All() = %d servers, %v", len(all), err) + } + if err := store.Remove(ScopeProject, "maestro"); err != nil { + t.Fatal(err) + } + if readJSON(t, filepath.Join(project, ".mcp.json"))["other"] == nil { + t.Fatal("remove must preserve other keys") + } + if err := store.Remove(ScopeUser, "missing"); err == nil { + t.Fatal("removing an unknown server should fail") + } +} + +func TestValidate(t *testing.T) { + for _, config := range []ServerConfig{{}, {Type: "http"}, {Type: "ws", URL: "https://x"}, {Command: "x", URL: "https://x"}} { + if config.Validate() == nil { + t.Errorf("expected invalid: %+v", config) + } + } + if err := (ServerConfig{Command: "npx"}).Validate(); err != nil { + t.Error(err) + } +} + +func TestExpandEnv(t *testing.T) { + lookup := func(name string) (string, bool) { + value, ok := map[string]string{"TOKEN": "abc", "EMPTY": ""}[name] + return value, ok + } + cases := map[string]string{ + "Bearer ${TOKEN}": "Bearer abc", + "${MISSING:-fallback}": "fallback", + "${EMPTY:-fallback}": "fallback", + "${MISSING}": "", + "$TOKEN stays": "$TOKEN stays", + "${TOKEN}/${MISSING:-x/y}": "abc/x/y", + } + for input, want := range cases { + if got := ExpandEnv(input, lookup); got != want { + t.Errorf("ExpandEnv(%q) = %q, want %q", input, got, want) + } + } + config := ServerConfig{Command: "${TOKEN}", Args: []string{"--k=${TOKEN}"}, Env: map[string]string{"A": "${TOKEN}"}, URL: "https://${MISSING:-host}/mcp", Headers: map[string]string{"Authorization": "Bearer ${TOKEN}"}} + expanded := config.Expanded(lookup) + if expanded.Command != "abc" || expanded.Args[0] != "--k=abc" || expanded.Env["A"] != "abc" || expanded.URL != "https://host/mcp" || expanded.Headers["Authorization"] != "Bearer abc" { + t.Fatalf("Expanded = %+v", expanded) + } + if config.Env["A"] != "${TOKEN}" { + t.Fatal("Expanded must not mutate the original config") + } +} + +func TestMasking(t *testing.T) { + config := ServerConfig{ + Command: "docker", Args: []string{"run", "--token=sk-live", "-e", "X"}, + Env: map[string]string{"SCALR_API_TOKEN": "secret"}, Headers: map[string]string{"Authorization": "Bearer secret"}, + URL: "https://user:pass@example.com/mcp?api_key=secret&x=1", + } + masked := config.Masked() + text, _ := json.Marshal(masked) + if strings.Contains(string(text), "secret") || strings.Contains(string(text), "pass@") || strings.Contains(string(text), "sk-live") { + t.Fatalf("masked config leaks a secret: %s", text) + } + pairs := MaskedPairs(masked.Env, "=") + if len(pairs) != 1 || pairs[0] != "SCALR_API_TOKEN="+Mask { + t.Fatalf("pairs = %v", pairs) + } +} + +func TestClientToolNamer(t *testing.T) { + namer := NewClientToolNamer() + if got := namer.Name("My Files", "read-file"); got != "local_my_files_read_file" { + t.Fatalf("got %q", got) + } + if got := namer.Name("my_files", "read file"); got != "local_my_files_read_file_2" { + t.Fatalf("collision: got %q", got) + } + long := namer.Name(strings.Repeat("server", 10), strings.Repeat("tool", 10)) + if len(long) > len(ClientToolPrefix)+60 { + t.Fatalf("name too long: %d", len(long)) + } + if got := namer.Name("Ünïcode!!", "✓"); !strings.HasPrefix(got, ClientToolPrefix) { + t.Fatalf("got %q", got) + } + for _, name := range []string{long, namer.Name("a", "b")} { + for _, r := range strings.TrimPrefix(name, ClientToolPrefix) { + if (r < 'a' || r > 'z') && (r < '0' || r > '9') && r != '_' { + t.Fatalf("invalid rune %q in %q", r, name) + } + } + } + again := namer.Name(strings.Repeat("server", 10), strings.Repeat("tool", 10)) + if again == long || len(again) > len(ClientToolPrefix)+60 { + t.Fatalf("truncated collision: %q vs %q", again, long) + } +} + +func TestImportedSourcesRealShapes(t *testing.T) { + home, project := t.TempDir(), t.TempDir() + writeFile(t, filepath.Join(home, ".claude.json"), marshalTestJSON(t, map[string]any{ + "numStartups": 3, + "mcpServers": map[string]any{ + "sentry": map[string]any{"type": "http", "url": "https://mcp.sentry.dev/mcp"}, + }, + "projects": map[string]any{ + project: map[string]any{ + "allowedTools": []string{}, + "mcpServers": map[string]any{ + "local-db": map[string]any{ + "type": "stdio", "command": "npx", "args": []string{"-y", "db"}, + "env": map[string]string{"DB": "x"}, + }, + }, + }, + }, + })) + writeFile(t, filepath.Join(home, ".codex", "config.toml"), `model = "gpt-5" +[mcp_servers.docs] +command = "npx" +args = ["-y", "docs-mcp"] +[mcp_servers.docs.env] +DOCS_TOKEN = "t" + +[mcp_servers.linear] +url = "https://mcp.linear.app/mcp" +bearer_token_env_var = "LINEAR_TOKEN" +`) + writeFile(t, filepath.Join(home, ".cursor", "mcp.json"), `{"mcpServers":{"cursor-user":{"command":"uvx","args":["x"]}}}`) + writeFile(t, filepath.Join(project, ".cursor", "mcp.json"), `{"mcpServers":{"cursor-project":{"url":"https://example.com/sse","type":"sse"}}}`) + writeFile(t, filepath.Join(project, ".vscode", "mcp.json"), `{"inputs":[],"servers":{"github":{"type":"http","url":"https://api.githubcopilot.com/mcp/"}}}`) + var desktop string + switch runtime.GOOS { + case "darwin": + desktop = filepath.Join(home, "Library", "Application Support", "Claude", "claude_desktop_config.json") + case "linux": + desktop = filepath.Join(home, ".config", "Claude", "claude_desktop_config.json") + } + if desktop != "" { + writeFile(t, desktop, `{"mcpServers":{"filesystem":{"command":"npx","args":["-y","@modelcontextprotocol/server-filesystem","/Users/me/Desktop"]}}}`) + } + + discovery := Discovery{Home: home, CWD: project, GOOS: runtime.GOOS} + servers, problems := discovery.Imported() + if len(problems) != 0 { + t.Fatalf("problems: %+v", problems) + } + got := map[string]Server{} + for _, server := range servers { + got[server.Name] = server + if server.Scope != ScopeImported || !server.ReadOnly() { + t.Errorf("%s should be imported read-only", server.Name) + } + } + expected := []string{"sentry", "local-db", "docs", "linear", "cursor-user", "cursor-project", "github"} + if desktop != "" { + expected = append(expected, "filesystem") + } + for _, name := range expected { + if _, ok := got[name]; !ok { + t.Errorf("missing imported server %s (got %v)", name, got) + } + } + if got["linear"].Config.Transport() != TransportHTTP || got["linear"].Config.BearerTokenEnvVar != "LINEAR_TOKEN" { + t.Errorf("codex url server: %+v", got["linear"].Config) + } + if got["docs"].Config.Env["DOCS_TOKEN"] != "t" || got["docs"].Config.Transport() != TransportStdio { + t.Errorf("codex stdio server: %+v", got["docs"].Config) + } + if got["cursor-project"].Config.Transport() != TransportSSE || got["github"].Config.Transport() != TransportHTTP { + t.Errorf("transports: %+v %+v", got["cursor-project"].Config, got["github"].Config) + } + + writeFile(t, filepath.Join(home, ".cursor", "mcp.json"), `{not json`) + if _, problems := discovery.Imported(); len(problems) != 1 { + t.Fatalf("broken file should be reported once, got %+v", problems) + } +} + +func TestStateApprovalAndEnable(t *testing.T) { + store := &StateStore{Path: filepath.Join(t.TempDir(), "mcp-state.json")} + project := Server{Name: "maestro", Scope: ScopeProject, Path: "/repo/.mcp.json", Config: ServerConfig{Command: "maestro"}} + imported := Server{Name: "sentry", Scope: ScopeImported, Path: "/home/.claude.json", Config: ServerConfig{Type: "http", URL: "https://x"}} + user := Server{Name: "fs", Scope: ScopeUser, Config: ServerConfig{Command: "npx"}} + + state, err := store.Load() + if err != nil { + t.Fatal(err) + } + if !state.NeedsApproval(project) || state.NeedsApproval(user) { + t.Fatal("only project servers need approval") + } + if state.IsEnabled(imported) || !state.IsEnabled(user) { + t.Fatal("imported servers start off, others on") + } + if err := store.Update(func(state *State) { state.Approve(project); state.SetEnabled(imported, true) }); err != nil { + t.Fatal(err) + } + state, _ = store.Load() + if state.NeedsApproval(project) || !state.IsEnabled(imported) { + t.Fatal("approval and enable should persist") + } + changed := project + changed.Config.Args = []string{"--evil"} + if !state.NeedsApproval(changed) { + t.Fatal("a config change must require approval again") + } + assertFileModePerm(t, store.Path, 0o600) +} + +func TestTruncateOutput(t *testing.T) { + short := "ok" + if TruncateOutput(short) != short { + t.Fatal("short output changed") + } + long := strings.Repeat("é", MaxToolOutputBytes) + got := TruncateOutput(long) + if len(got) > MaxToolOutputBytes || !strings.Contains(got, "truncated") { + t.Fatalf("len=%d", len(got)) + } +} + +func TestTokenStorePerms(t *testing.T) { + store := &TokenStore{Path: filepath.Join(t.TempDir(), "mcp-oauth.json")} + if err := store.Put("https://api.braintrust.dev/mcp", StoredCredential{ClientID: "c", AccessToken: "a"}); err != nil { + t.Fatal(err) + } + assertFileModePerm(t, store.Path, 0o600) + if credential, ok := store.Get("https://api.braintrust.dev/mcp"); !ok || credential.AccessToken != "a" { + t.Fatal("token not stored") + } + if removed, err := store.Delete("https://api.braintrust.dev/mcp"); err != nil || !removed { + t.Fatal("token not deleted") + } +} + +func TestManagerLoadReportsProblemsAndRevokesApproval(t *testing.T) { + home, project := t.TempDir(), t.TempDir() + store := &Store{UserPath: filepath.Join(home, ".rootly-cli", "mcp.json"), CWD: project} + states := &StateStore{Path: filepath.Join(home, ".rootly-cli", "mcp-state.json")} + writeFile(t, filepath.Join(project, ".mcp.json"), `{"mcpServers":{"maestro":{"command":"maestro-not-installed"}}}`) + writeFile(t, filepath.Join(home, ".cursor", "mcp.json"), `{not json`) + manager := NewManager(ManagerOptions{CWD: project, Store: store, State: states, Discovery: &Discovery{Home: home, CWD: project}}) + defer manager.Close() + + if err := manager.Load(); err == nil || !strings.Contains(err.Error(), "mcp.json") { + t.Fatalf("broken import should be reported, got %v", err) + } + if len(manager.Problems()) == 0 { + t.Fatal("expected problems to be kept for /mcp") + } + views := manager.Views() + if len(views) != 1 || views[0].Status != StatusPending { + t.Fatalf("project server should load as pending, got %+v", views) + } + if err := manager.Approve(views[0].Server.ID()); err != nil { + t.Fatal(err) + } + if err := states.Update(func(state *State) { state.Approvals = nil }); err != nil { + t.Fatal(err) + } + _ = manager.Load() + if got := manager.Views()[0].Status; got != StatusPending { + t.Fatalf("revoked approval should drop back to pending, got %s", got) + } +} diff --git a/internal/mcp/names.go b/internal/mcp/names.go new file mode 100644 index 0000000..9d422e7 --- /dev/null +++ b/internal/mcp/names.go @@ -0,0 +1,75 @@ +package mcp + +import ( + "fmt" + "strings" +) + +const ( + ClientToolPrefix = "local_" + maxClientToolName = 60 + maxServerNamePart = 20 +) + +func sanitizeNamePart(value string) string { + var builder strings.Builder + lastUnderscore := false + for _, r := range strings.ToLower(value) { + if (r >= 'a' && r <= 'z') || (r >= '0' && r <= '9') { + builder.WriteRune(r) + lastUnderscore = false + continue + } + if !lastUnderscore { + builder.WriteByte('_') + lastUnderscore = true + } + } + return strings.Trim(builder.String(), "_") +} + +// ClientToolNamer builds unique `local__` names matching +// ^local_[a-z0-9_]{1,60}$ and remembers the original server/tool. +type ClientToolNamer struct { + used map[string]bool +} + +func NewClientToolNamer() *ClientToolNamer { + return &ClientToolNamer{used: map[string]bool{}} +} + +func (n *ClientToolNamer) Reserve(name string) { + name = strings.TrimPrefix(name, ClientToolPrefix) + if name != "" { + n.used[name] = true + } +} + +func (n *ClientToolNamer) Name(server, tool string) string { + serverPart := sanitizeNamePart(server) + if len(serverPart) > maxServerNamePart { + serverPart = strings.TrimRight(serverPart[:maxServerNamePart], "_") + } + if serverPart == "" { + serverPart = "server" + } + toolPart := sanitizeNamePart(tool) + if toolPart == "" { + toolPart = "tool" + } + body := serverPart + "_" + toolPart + if len(body) > maxClientToolName { + body = strings.TrimRight(body[:maxClientToolName], "_") + } + candidate := body + for suffix := 2; n.used[candidate]; suffix++ { + tail := fmt.Sprintf("_%d", suffix) + trimmed := body + if len(trimmed)+len(tail) > maxClientToolName { + trimmed = strings.TrimRight(trimmed[:maxClientToolName-len(tail)], "_") + } + candidate = trimmed + tail + } + n.used[candidate] = true + return ClientToolPrefix + candidate +} diff --git a/internal/mcp/oauth.go b/internal/mcp/oauth.go new file mode 100644 index 0000000..fbcae4b --- /dev/null +++ b/internal/mcp/oauth.go @@ -0,0 +1,342 @@ +package mcp + +import ( + "context" + "encoding/json" + "errors" + "fmt" + "html" + "io" + "net" + "net/http" + "os" + "os/exec" + "path/filepath" + "runtime" + "strings" + "sync" + "time" + + "github.com/modelcontextprotocol/go-sdk/auth" + "github.com/modelcontextprotocol/go-sdk/oauthex" + "golang.org/x/oauth2" + + "github.com/rootlyhq/rootly-cli/internal/config" +) + +const tokenFile = "mcp-oauth.json" + +// ErrNeedsAuth means an HTTP server asked for OAuth and no usable token is stored. +var ErrNeedsAuth = errors.New("needs authentication") + +type StoredCredential struct { + ClientID string `json:"client_id"` + ClientSecret string `json:"client_secret,omitempty"` + AuthURL string `json:"auth_url"` + TokenURL string `json:"token_url"` + AuthStyle int `json:"auth_style,omitempty"` + Scopes []string `json:"scopes,omitempty"` + RedirectURL string `json:"redirect_url,omitempty"` + AccessToken string `json:"access_token"` + RefreshToken string `json:"refresh_token,omitempty"` + TokenType string `json:"token_type,omitempty"` + Expiry time.Time `json:"expiry,omitempty"` +} + +func (c StoredCredential) oauthConfig() *oauth2.Config { + return &oauth2.Config{ + ClientID: c.ClientID, + ClientSecret: c.ClientSecret, + Endpoint: oauth2.Endpoint{AuthURL: c.AuthURL, TokenURL: c.TokenURL, AuthStyle: oauth2.AuthStyle(c.AuthStyle)}, + Scopes: c.Scopes, + RedirectURL: c.RedirectURL, + } +} + +func (c StoredCredential) token() *oauth2.Token { + return &oauth2.Token{AccessToken: c.AccessToken, RefreshToken: c.RefreshToken, TokenType: c.TokenType, Expiry: c.Expiry} +} + +// TokenStore keeps MCP OAuth credentials keyed by server URL in a 0600 file. +type TokenStore struct { + Path string + mu sync.Mutex +} + +func NewTokenStore() *TokenStore { + return &TokenStore{Path: filepath.Join(config.Dir(), tokenFile)} +} + +func (s *TokenStore) readAll() (map[string]StoredCredential, error) { + credentials := map[string]StoredCredential{} + data, err := os.ReadFile(s.Path) + if err != nil { + if os.IsNotExist(err) { + return credentials, nil + } + return nil, err + } + if len(data) == 0 { + return credentials, nil + } + if err := json.Unmarshal(data, &credentials); err != nil { + return nil, fmt.Errorf("parse %s: %w", s.Path, err) + } + return credentials, nil +} + +func (s *TokenStore) writeAll(credentials map[string]StoredCredential) error { + document := jsonDocument{} + for key, credential := range credentials { + encoded, err := json.Marshal(credential) + if err != nil { + return err + } + document[key] = encoded + } + if err := writeJSONDocument(s.Path, document, 0o600); err != nil { + return err + } + return os.Chmod(s.Path, 0o600) +} + +func (s *TokenStore) Get(serverURL string) (StoredCredential, bool) { + s.mu.Lock() + defer s.mu.Unlock() + credentials, err := s.readAll() + if err != nil { + return StoredCredential{}, false + } + credential, ok := credentials[tokenKey(serverURL)] + return credential, ok && credential.AccessToken != "" +} + +func (s *TokenStore) Put(serverURL string, credential StoredCredential) error { + s.mu.Lock() + defer s.mu.Unlock() + credentials, err := s.readAll() + if err != nil { + return err + } + credentials[tokenKey(serverURL)] = credential + return s.writeAll(credentials) +} + +func (s *TokenStore) Delete(serverURL string) (bool, error) { + s.mu.Lock() + defer s.mu.Unlock() + credentials, err := s.readAll() + if err != nil { + return false, err + } + key := tokenKey(serverURL) + if _, ok := credentials[key]; !ok { + return false, nil + } + delete(credentials, key) + return true, s.writeAll(credentials) +} + +func tokenKey(serverURL string) string { + return strings.TrimRight(strings.TrimSpace(serverURL), "/") +} + +type persistingTokenSource struct { + store *TokenStore + serverURL string + base StoredCredential + source oauth2.TokenSource + mu sync.Mutex + last string +} + +func (p *persistingTokenSource) Token() (*oauth2.Token, error) { + token, err := p.source.Token() + if err != nil { + return nil, err + } + p.mu.Lock() + defer p.mu.Unlock() + if token.AccessToken != p.last { + p.last = token.AccessToken + credential := p.base + credential.AccessToken = token.AccessToken + credential.TokenType = token.TokenType + credential.Expiry = token.Expiry + if token.RefreshToken != "" { + credential.RefreshToken = token.RefreshToken + } + p.base = credential + _ = p.store.Put(p.serverURL, credential) + } + return token, nil +} + +func (s *TokenStore) tokenSource(serverURL string, credential StoredCredential, cfg *oauth2.Config, token *oauth2.Token, client *http.Client) oauth2.TokenSource { + return s.tokenSourceWithContext(context.Background(), serverURL, credential, cfg, token, client) +} + +func (s *TokenStore) tokenSourceWithContext(ctx context.Context, serverURL string, credential StoredCredential, cfg *oauth2.Config, token *oauth2.Token, client *http.Client) oauth2.TokenSource { + if ctx == nil { + ctx = context.Background() + } + if client != nil { + ctx = context.WithValue(ctx, oauth2.HTTPClient, client) + } + return &persistingTokenSource{store: s, serverURL: serverURL, base: credential, source: cfg.TokenSource(ctx, token), last: token.AccessToken} +} + +type LoginOptions struct { + // OpenURL shows the authorization URL to the user (browser + print). + OpenURL func(string) + Client *http.Client +} + +// Login runs the MCP authorization flow (protected-resource discovery, +// dynamic client registration, PKCE, loopback redirect) and stores the token. +func (s *TokenStore) Login(ctx context.Context, server ServerConfig, options LoginOptions) error { + if server.Transport() == TransportStdio { + return errors.New("OAuth login only applies to http and sse servers") + } + serverURL := server.URL + listener, err := (&net.ListenConfig{}).Listen(ctx, "tcp", "127.0.0.1:0") + if err != nil { + return fmt.Errorf("start OAuth callback listener: %w", err) + } + defer func() { _ = listener.Close() }() + redirectURL := fmt.Sprintf("http://127.0.0.1:%d/callback", listener.Addr().(*net.TCPAddr).Port) + + results := make(chan *auth.AuthorizationResult, 1) + failures := make(chan error, 1) + httpServer := &http.Server{ReadHeaderTimeout: 10 * time.Second, Handler: http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + if r.URL.Path != "/callback" { + http.NotFound(w, r) + return + } + query := r.URL.Query() + if problem := query.Get("error"); problem != "" { + w.Header().Set("Content-Type", "text/html; charset=utf-8") + _, _ = fmt.Fprintf(w, "

Authorization failed

%s

", html.EscapeString(problem)) + select { + case failures <- fmt.Errorf("authorization failed: %s", problem): + default: + } + return + } + w.Header().Set("Content-Type", "text/html; charset=utf-8") + _, _ = io.WriteString(w, "

Rootly CLI is connected.

You can close this tab.

") + select { + case results <- &auth.AuthorizationResult{Code: query.Get("code"), State: query.Get("state"), Iss: query.Get("iss")}: + default: + } + })} + go func() { _ = httpServer.Serve(listener) }() + defer func() { _ = httpServer.Close() }() + + client := options.Client + if client == nil { + client = http.DefaultClient + } + var saveErr error + handler, err := auth.NewAuthorizationCodeHandler(&auth.AuthorizationCodeHandlerConfig{ + DynamicClientRegistrationConfig: &auth.DynamicClientRegistrationConfig{Metadata: &oauthex.ClientRegistrationMetadata{ + ClientName: "Rootly CLI", + RedirectURIs: []string{redirectURL}, + GrantTypes: []string{"authorization_code", "refresh_token"}, + ResponseTypes: []string{"code"}, + TokenEndpointAuthMethod: "none", + }}, + RedirectURL: redirectURL, + RequestRefreshToken: true, + Client: client, + AuthorizationCodeFetcher: func(ctx context.Context, args *auth.AuthorizationArgs) (*auth.AuthorizationResult, error) { + if options.OpenURL != nil { + options.OpenURL(args.URL) + } + select { + case result := <-results: + return result, nil + case err := <-failures: + return nil, err + case <-ctx.Done(): + return nil, ctx.Err() + } + }, + NewTokenSource: func(_ context.Context, cfg *oauth2.Config, token *oauth2.Token) (oauth2.TokenSource, error) { + credential := StoredCredential{ + ClientID: cfg.ClientID, ClientSecret: cfg.ClientSecret, + AuthURL: cfg.Endpoint.AuthURL, TokenURL: cfg.Endpoint.TokenURL, AuthStyle: int(cfg.Endpoint.AuthStyle), + Scopes: cfg.Scopes, RedirectURL: cfg.RedirectURL, + AccessToken: token.AccessToken, RefreshToken: token.RefreshToken, TokenType: token.TokenType, Expiry: token.Expiry, + } + saveErr = s.Put(serverURL, credential) + return s.tokenSource(serverURL, credential, cfg, token, client), nil + }, + }) + if err != nil { + return err + } + + request, response, err := probe(ctx, client, server) + if err != nil { + return err + } + if response.StatusCode != http.StatusUnauthorized && response.StatusCode != http.StatusForbidden { + _ = response.Body.Close() + return fmt.Errorf("%s did not ask for authorization (HTTP %d); no login needed", serverURL, response.StatusCode) + } + if err := handler.Authorize(ctx, request, response); err != nil { + return err + } + return saveErr +} + +func probe(ctx context.Context, client *http.Client, server ServerConfig) (*http.Request, *http.Response, error) { + method, body := http.MethodPost, `{"jsonrpc":"2.0","id":0,"method":"ping"}` + if server.Transport() == TransportSSE { + method, body = http.MethodGet, "" + } + request, err := http.NewRequestWithContext(ctx, method, server.URL, strings.NewReader(body)) + if err != nil { + return nil, nil, err + } + request.Header.Set("Accept", "application/json, text/event-stream") + if body != "" { + request.Header.Set("Content-Type", "application/json") + } + for key, value := range server.Headers { + request.Header.Set(key, value) + } + guarded := *client + guarded.CheckRedirect = rejectCrossOriginRedirect + response, err := guarded.Do(request) + if err != nil { + return nil, nil, err + } + return request, response, nil +} + +func OpenBrowser(url string) error { + switch runtime.GOOS { + case "darwin": + return exec.CommandContext(context.Background(), "open", url).Start() + case "windows": + return exec.CommandContext(context.Background(), "rundll32", "url.dll,FileProtocolHandler", url).Start() + default: + return exec.CommandContext(context.Background(), "xdg-open", url).Start() + } +} + +// ProbeStatus sends one unauthenticated MCP request and returns the HTTP status. +func ProbeStatus(ctx context.Context, client *http.Client, server ServerConfig) (int, error) { + if client == nil { + client = http.DefaultClient + } + _, response, err := probe(ctx, client, server) + if err != nil { + return 0, err + } + _, _ = io.Copy(io.Discard, io.LimitReader(response.Body, 64*1024)) + _ = response.Body.Close() + return response.StatusCode, nil +} diff --git a/internal/mcp/registry.go b/internal/mcp/registry.go new file mode 100644 index 0000000..1363312 --- /dev/null +++ b/internal/mcp/registry.go @@ -0,0 +1,433 @@ +package mcp + +import ( + "context" + "encoding/json" + "errors" + "fmt" + "net/http" + "net/url" + "os" + "path/filepath" + "regexp" + "sort" + "strings" + "time" + + "github.com/rootlyhq/rootly-cli/internal/config" +) + +const ( + RegistryURL = "https://registry.modelcontextprotocol.io/v0/servers" + indexFile = "mcp-index.json" + registryPageLimit = 100 + registryMaxPages = 200 + IndexMaxAge = 7 * 24 * time.Hour + officialMetaKey = "io.modelcontextprotocol.registry/official" + maxMarketNameBytes = 40 +) + +// MarketEntry is one installable server from the curated list or the official registry. +type MarketEntry struct { + Name string `json:"name"` + RegistryName string `json:"registry_name,omitempty"` + Title string `json:"title,omitempty"` + Description string `json:"description,omitempty"` + Author string `json:"author,omitempty"` + EnvHint string `json:"env_hint,omitempty"` + Rootly bool `json:"rootly,omitempty"` + Config ServerConfig `json:"config"` +} + +type marketIndex struct { + IndexedAt time.Time `json:"indexed_at"` + Entries []MarketEntry `json:"entries"` +} + +func npxConfig(pkg string, args ...string) ServerConfig { + return ServerConfig{Command: "npx", Args: append([]string{"-y", pkg}, args...)} +} + +func uvxConfig(pkg string) ServerConfig { return ServerConfig{Command: "uvx", Args: []string{pkg}} } + +// CuratedEntries lists vendors Rootly integrates with first, then common servers. +func CuratedEntries() []MarketEntry { + return []MarketEntry{ + {Name: "datadog", Author: "datadog", Rootly: true, Description: "Datadog: monitors, metrics, logs (remote, OAuth).", Config: ServerConfig{Type: TransportHTTP, URL: "https://mcp.datadoghq.com/api/unstable/mcp-server/mcp"}}, + {Name: "sentry", Author: "sentry", Rootly: true, Description: "Sentry: issues and events (remote, OAuth).", Config: ServerConfig{Type: TransportHTTP, URL: "https://mcp.sentry.dev/mcp"}}, + {Name: "linear", Author: "linear", Rootly: true, Description: "Linear: issues and projects (remote, OAuth).", Config: ServerConfig{Type: TransportHTTP, URL: "https://mcp.linear.app/mcp"}}, + {Name: "notion", Author: "notion", Rootly: true, Description: "Notion: pages and databases (remote, OAuth).", Config: ServerConfig{Type: TransportHTTP, URL: "https://mcp.notion.com/mcp"}}, + {Name: "github", Author: "modelcontextprotocol", Description: "GitHub: repos, issues, PRs, code search.", EnvHint: "GITHUB_PERSONAL_ACCESS_TOKEN", Config: npxConfig("@modelcontextprotocol/server-github")}, + {Name: "gitlab", Author: "modelcontextprotocol", Description: "GitLab: projects, issues, MRs.", EnvHint: "GITLAB_PERSONAL_ACCESS_TOKEN", Config: npxConfig("@modelcontextprotocol/server-gitlab")}, + {Name: "filesystem", Author: "modelcontextprotocol", Description: "Local filesystem access, scoped to a path.", Config: npxConfig("@modelcontextprotocol/server-filesystem", "${PWD}")}, + {Name: "postgres", Author: "modelcontextprotocol", Description: "Query a PostgreSQL database (read-only).", EnvHint: "DATABASE_URL", Config: npxConfig("@modelcontextprotocol/server-postgres", "${DATABASE_URL}")}, + {Name: "slack", Author: "modelcontextprotocol", Description: "Slack: channels, messages, users.", EnvHint: "SLACK_BOT_TOKEN", Config: npxConfig("@modelcontextprotocol/server-slack")}, + {Name: "memory", Author: "modelcontextprotocol", Description: "Persistent knowledge-graph memory.", Config: npxConfig("@modelcontextprotocol/server-memory")}, + {Name: "fetch", Author: "modelcontextprotocol", Description: "Fetch web pages as markdown.", Config: uvxConfig("mcp-server-fetch")}, + {Name: "aws-api", Author: "awslabs", Description: "AWS: call AWS APIs with your local credentials.", EnvHint: "AWS_PROFILE", Config: uvxConfig("awslabs.aws-api-mcp-server@latest")}, + {Name: "aws-docs", Author: "awslabs", Description: "AWS documentation: search and read.", Config: uvxConfig("awslabs.aws-documentation-mcp-server@latest")}, + {Name: "aws-cloudwatch", Author: "awslabs", Description: "AWS CloudWatch: metrics, logs, alarms.", EnvHint: "AWS_PROFILE", Config: uvxConfig("awslabs.cloudwatch-mcp-server@latest")}, + } +} + +type registryPage struct { + Servers []struct { + Server registryServer `json:"server"` + Meta map[string]json.RawMessage `json:"_meta"` + } `json:"servers"` + Metadata struct { + NextCursor string `json:"nextCursor"` + } `json:"metadata"` +} + +type registryServer struct { + Name string `json:"name"` + Title string `json:"title"` + Description string `json:"description"` + Repository struct { + URL string `json:"url"` + } `json:"repository"` + Packages []struct { + RegistryType string `json:"registryType"` + Identifier string `json:"identifier"` + Transport struct { + Type string `json:"type"` + } `json:"transport"` + } `json:"packages"` + Remotes []struct { + Type string `json:"type"` + URL string `json:"url"` + } `json:"remotes"` +} + +var nonSlug = regexp.MustCompile(`[^a-z0-9]+`) + +func marketName(registryName string) string { + last := registryName[strings.LastIndex(registryName, "/")+1:] + name := strings.Trim(nonSlug.ReplaceAllString(strings.ToLower(last), "-"), "-") + if len(name) > maxMarketNameBytes { + name = strings.Trim(name[:maxMarketNameBytes], "-") + } + return name +} + +var githubOwner = regexp.MustCompile(`(?i)github\.com/([^/]+)`) + +func registryAuthor(server registryServer) string { + if match := githubOwner.FindStringSubmatch(server.Repository.URL); match != nil { + return match[1] + } + parts := strings.Split(strings.SplitN(server.Name, "/", 2)[0], ".") + if len(parts) > 2 && parts[0] == "io" && parts[1] == "github" { + return parts[2] + } + if len(parts) > 1 { + return parts[1] + } + return parts[0] +} + +// installFromRegistry maps npm→npx -y, pypi→uvx, oci→docker run -i --rm, else the first remote URL. +func installFromRegistry(server registryServer) (ServerConfig, bool) { + for _, pkg := range server.Packages { + if pkg.Identifier == "" || (pkg.Transport.Type != "" && pkg.Transport.Type != "stdio") { + continue + } + switch pkg.RegistryType { + case "npm": + return npxConfig(pkg.Identifier), true + case "pypi": + return uvxConfig(pkg.Identifier), true + case "oci": + return ServerConfig{Command: "docker", Args: []string{"run", "-i", "--rm", pkg.Identifier}}, true + } + } + for _, remote := range server.Remotes { + if remote.URL == "" { + continue + } + if remote.Type == "sse" { + return ServerConfig{Type: TransportSSE, URL: remote.URL}, true + } + return ServerConfig{Type: TransportHTTP, URL: remote.URL}, true + } + return ServerConfig{}, false +} + +func isLatest(meta map[string]json.RawMessage) bool { + raw, ok := meta[officialMetaKey] + if !ok { + return true + } + var official struct { + IsLatest bool `json:"isLatest"` + } + return json.Unmarshal(raw, &official) == nil && official.IsLatest +} + +// FetchRegistry pages through the official registry and keeps the latest version of each installable server. +func FetchRegistry(ctx context.Context, client *http.Client, baseURL string, progress func(entries, pages int)) ([]MarketEntry, error) { + if client == nil { + client = &http.Client{Timeout: 20 * time.Second} + } + var entries []MarketEntry + seen := map[string]bool{} + cursor := "" + for page := 1; page <= registryMaxPages; page++ { + next, err := fetchRegistryPage(ctx, client, baseURL, cursor) + if err != nil { + if len(entries) > 0 { + return entries, fmt.Errorf("registry page %d: %w (index incomplete)", page, err) + } + return nil, err + } + for _, item := range next.Servers { + if item.Server.Name == "" || !isLatest(item.Meta) || seen[item.Server.Name] { + continue + } + install, ok := installFromRegistry(item.Server) + name := marketName(item.Server.Name) + if !ok || name == "" || install.Validate() != nil { + continue + } + seen[item.Server.Name] = true + entries = append(entries, MarketEntry{Name: name, RegistryName: item.Server.Name, Title: item.Server.Title, Description: item.Server.Description, Author: registryAuthor(item.Server), Config: install}) + } + if progress != nil { + progress(len(entries), page) + } + if next.Metadata.NextCursor == "" { + break + } + cursor = next.Metadata.NextCursor + } + return entries, nil +} + +func fetchRegistryPage(ctx context.Context, client *http.Client, baseURL, cursor string) (*registryPage, error) { + endpoint, err := url.Parse(baseURL) + if err != nil { + return nil, err + } + query := endpoint.Query() + query.Set("limit", fmt.Sprint(registryPageLimit)) + if cursor != "" { + query.Set("cursor", cursor) + } + endpoint.RawQuery = query.Encode() + request, err := http.NewRequestWithContext(ctx, http.MethodGet, endpoint.String(), http.NoBody) + if err != nil { + return nil, err + } + request.Header.Set("Accept", "application/json") + response, err := client.Do(request) + if err != nil { + return nil, err + } + defer func() { _ = response.Body.Close() }() + if response.StatusCode != http.StatusOK { + return nil, fmt.Errorf("registry returned HTTP %d", response.StatusCode) + } + var page registryPage + if err := json.NewDecoder(response.Body).Decode(&page); err != nil { + return nil, fmt.Errorf("decode registry page: %w", err) + } + return &page, nil +} + +func IndexPath() string { return filepath.Join(config.Dir(), indexFile) } + +func SaveIndex(path string, entries []MarketEntry, now time.Time) error { + data, err := json.Marshal(marketIndex{IndexedAt: now, Entries: entries}) + if err != nil { + return err + } + if err := os.MkdirAll(filepath.Dir(path), 0o700); err != nil { + return err + } + return os.WriteFile(path, data, 0o600) +} + +func LoadIndex(path string) (entries []MarketEntry, indexedAt time.Time, err error) { + data, err := os.ReadFile(path) + if err != nil { + return nil, time.Time{}, err + } + var index marketIndex + if err := json.Unmarshal(data, &index); err != nil { + return nil, time.Time{}, err + } + return index.Entries, index.IndexedAt, nil +} + +// Catalog merges curated entries (first) with the registry index. A curated entry +// replaces registry entries of the same name; registry entries that share a short +// name get distinct install names so none of them disappear. +func Catalog(indexed []MarketEntry) []MarketEntry { + curated := CuratedEntries() + seen := map[string]bool{} + curatedInstalls := map[string]bool{} + entries := make([]MarketEntry, 0, len(curated)+len(indexed)) + for _, entry := range curated { + seen[entry.Name] = true + curatedInstalls[entry.Name+"\x00"+installKey(entry.Config)] = true + entries = append(entries, entry) + } + for _, entry := range indexed { + if entry.Name == "" || curatedInstalls[entry.Name+"\x00"+installKey(entry.Config)] { + continue + } + if seen[entry.Name] { + entry.Name = distinctMarketName(entry, seen) + } + seen[entry.Name] = true + entries = append(entries, entry) + } + return entries +} + +func installKey(server ServerConfig) string { + return strings.Join(append([]string{server.Transport(), server.URL, server.Command}, server.Args...), "\x00") +} + +func distinctMarketName(entry MarketEntry, seen map[string]bool) string { + base := entry.Name + if publisher, _, ok := strings.Cut(entry.RegistryName, "/"); ok { + if owner := marketName(publisher[strings.LastIndex(publisher, ".")+1:]); owner != "" { + base = marketName(entry.Name + "-" + owner) + } + } + name := base + for suffix := 2; seen[name]; suffix++ { + tail := fmt.Sprintf("-%d", suffix) + name = strings.Trim(base[:min(len(base), maxMarketNameBytes-len(tail))], "-") + tail + } + return name +} + +func fuzzyScore(query, text string) int { + text = strings.ToLower(text) + if index := strings.Index(text, query); index >= 0 { + return 100 - min(index, 99) + } + matched := 0 + for _, char := range text { + if matched < len(query) && byte(char) == query[matched] { + matched++ + } + } + if matched == len(query) { + return 1 + } + return -1 +} + +// SearchMarket ranks entries by name, then title and description; ties keep catalog order. +func SearchMarket(entries []MarketEntry, query string, limit int) []MarketEntry { + query = strings.ToLower(strings.TrimSpace(query)) + if query == "" { + return entries[:min(limit, len(entries))] + } + type scored struct { + entry MarketEntry + score int + } + var matches []scored + for _, entry := range entries { + score := max(fuzzyScore(query, entry.Name)*2, fuzzyScore(query, entry.Title+" "+entry.Description)) + if score >= 0 { + matches = append(matches, scored{entry, score}) + } + } + sort.SliceStable(matches, func(i, j int) bool { return matches[i].score > matches[j].score }) + results := make([]MarketEntry, 0, min(limit, len(matches))) + for _, match := range matches[:min(limit, len(matches))] { + results = append(results, match.entry) + } + return results +} + +// InstallLabel is the exact command or URL that adding the entry will run or contact. +func InstallLabel(entry MarketEntry) string { + if entry.Config.URL != "" { + return entry.Config.Type + " " + entry.Config.URL + } + return strings.TrimSpace(entry.Config.Command + " " + strings.Join(entry.Config.Args, " ")) +} + +// AddFromMarket writes the entry to USER scope switched off, so nothing runs until the user turns it on. +func AddFromMarket(store *Store, states *StateStore, entry MarketEntry) (Server, error) { + if err := ValidateServerName(entry.Name); err != nil { + return Server{}, err + } + if err := entry.Config.Validate(); err != nil { + return Server{}, err + } + existing, err := store.All() + if err != nil { + return Server{}, err + } + for _, server := range existing { + if server.Name == entry.Name { + return Server{}, fmt.Errorf("an MCP server named %s already exists (%s scope)", entry.Name, server.Scope) + } + } + return AddDisabledUserServer(store, states, entry.Name, entry.Config) +} + +// AddDisabledUserServer stores a user server with its enabled state turned off. +func AddDisabledUserServer(store *Store, states *StateStore, name string, serverConfig ServerConfig) (Server, error) { + server := Server{Name: name, Scope: ScopeUser, Path: store.UserPath, Config: serverConfig} + if err := states.Update(func(state *State) { state.SetEnabled(server, false) }); err != nil { + return Server{}, err + } + if err := store.Add(ScopeUser, name, serverConfig); err != nil { + return Server{}, errors.Join(err, states.Update(func(state *State) { state.SetEnabled(server, true) })) + } + return server, nil +} + +// LoadCatalog uses the cached index, refetching it when missing, stale or forced; on fetch failure it falls back to the cache or the curated list. +func LoadCatalog(ctx context.Context, client *http.Client, registryURL, indexPath string, refresh bool, progress func(entries, pages int)) ([]MarketEntry, error) { + cached, indexedAt, cacheErr := LoadIndex(indexPath) + if !refresh && cacheErr == nil && time.Since(indexedAt) < IndexMaxAge { + return Catalog(cached), nil + } + fetched, err := FetchRegistry(ctx, client, registryURL, progress) + if err != nil { + if cacheErr != nil && len(fetched) > 0 { + cached = fetched + } + return Catalog(cached), fmt.Errorf("could not refresh the MCP registry index: %w", err) + } + if err := SaveIndex(indexPath, fetched, time.Now()); err != nil { + return Catalog(fetched), fmt.Errorf("could not cache the MCP registry index: %w", err) + } + return Catalog(fetched), nil +} + +// AddCommand is the equivalent `rootly mcp add` invocation for an entry, shell-quoted +// so `${VAR}` placeholders and URL query strings survive a copy-paste. +func AddCommand(entry MarketEntry) string { + parts := []string{"rootly", "mcp", "add", "-s", "user"} + if entry.Config.URL != "" { + parts = append(parts, "-t", entry.Config.Type, entry.Name, entry.Config.URL) + } else { + parts = append(parts, entry.Name, "--", entry.Config.Command) + parts = append(parts, entry.Config.Args...) + } + quoted := make([]string, 0, len(parts)) + for _, part := range parts { + quoted = append(quoted, shellQuote(part)) + } + return strings.Join(quoted, " ") +} + +var shellSafe = regexp.MustCompile(`^[A-Za-z0-9_@%+=:,./-]+$`) + +func shellQuote(value string) string { + if shellSafe.MatchString(value) { + return value + } + return "'" + strings.ReplaceAll(value, "'", `'\''`) + "'" +} diff --git a/internal/mcp/registry_test.go b/internal/mcp/registry_test.go new file mode 100644 index 0000000..b4bd262 --- /dev/null +++ b/internal/mcp/registry_test.go @@ -0,0 +1,225 @@ +package mcp + +import ( + "context" + "net/http" + "net/http/httptest" + "os" + "path/filepath" + "strings" + "testing" + "time" +) + +const registryPage1 = `{"servers":[ + {"server":{"name":"io.github.acme/pager-mcp","description":"Pager","packages":[{"registryType":"npm","identifier":"@acme/pager-mcp","transport":{"type":"stdio"}}]},"_meta":{"io.modelcontextprotocol.registry/official":{"isLatest":false}}}, + {"server":{"name":"io.github.acme/pager-mcp","description":"Pager latest","packages":[{"registryType":"npm","identifier":"@acme/pager-mcp","transport":{"type":"stdio"}}]},"_meta":{"io.modelcontextprotocol.registry/official":{"isLatest":true}}}, + {"server":{"name":"com.example/py-tool","description":"Py","packages":[{"registryType":"pypi","identifier":"py-tool","transport":{"type":"stdio"}}]},"_meta":{"io.modelcontextprotocol.registry/official":{"isLatest":true}}} +],"metadata":{"nextCursor":"c2"}}` + +const registryPage2 = `{"servers":[ + {"server":{"name":"io.github.box/oci-tool","repository":{"url":"https://github.com/boxco/oci"},"packages":[{"registryType":"oci","identifier":"ghcr.io/box/oci:1","transport":{"type":"stdio"}}]},"_meta":{"io.modelcontextprotocol.registry/official":{"isLatest":true}}}, + {"server":{"name":"dev.remote/events","remotes":[{"type":"sse","url":"https://events.example/sse"}]},"_meta":{"io.modelcontextprotocol.registry/official":{"isLatest":true}}}, + {"server":{"name":"dev.remote/api","remotes":[{"type":"streamable-http","url":"https://api.example/mcp"}]},"_meta":{"io.modelcontextprotocol.registry/official":{"isLatest":true}}}, + {"server":{"name":"dev.empty/nothing"},"_meta":{"io.modelcontextprotocol.registry/official":{"isLatest":true}}} +],"metadata":{}}` + +func TestFetchRegistryPagesAndMapsInstalls(t *testing.T) { + var queries []string + server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + queries = append(queries, r.URL.RawQuery) + if r.URL.Query().Get("cursor") == "c2" { + _, _ = w.Write([]byte(registryPage2)) + return + } + _, _ = w.Write([]byte(registryPage1)) + })) + defer server.Close() + entries, err := FetchRegistry(context.Background(), server.Client(), server.URL+"/v0/servers", nil) + if err != nil { + t.Fatal(err) + } + if len(queries) != 2 || queries[0] != "limit=100" || queries[1] != "cursor=c2&limit=100" { + t.Fatalf("queries = %v", queries) + } + got := map[string]string{} + for _, entry := range entries { + got[entry.Name] = InstallLabel(entry) + } + want := map[string]string{ + "pager-mcp": "npx -y @acme/pager-mcp", + "py-tool": "uvx py-tool", + "oci-tool": "docker run -i --rm ghcr.io/box/oci:1", + "events": "sse https://events.example/sse", + "api": "http https://api.example/mcp", + } + if len(got) != len(want) { + t.Fatalf("entries = %v", got) + } + for name, label := range want { + if got[name] != label { + t.Fatalf("%s = %q, want %q", name, got[name], label) + } + } + for _, entry := range entries { + if entry.Name == "pager-mcp" && entry.Description != "Pager latest" { + t.Fatalf("kept a non-latest version: %+v", entry) + } + if entry.Name == "oci-tool" && entry.Author != "boxco" { + t.Fatalf("author = %q", entry.Author) + } + } +} + +func TestCatalogPutsRootlyVendorsFirstAndSearchRanks(t *testing.T) { + catalog := Catalog([]MarketEntry{{Name: "sentry", Description: "registry duplicate", Config: ServerConfig{Type: TransportHTTP, URL: "https://mcp.sentry.dev/mcp"}}, {Name: "pager-mcp", Description: "Paging tool", Config: npxConfig("@acme/pager-mcp")}}) + var first []string + for _, entry := range catalog[:4] { + first = append(first, entry.Name) + } + if strings.Join(first, ",") != "datadog,sentry,linear,notion" { + t.Fatalf("first = %v", first) + } + for _, entry := range catalog { + if entry.Description == "registry duplicate" { + t.Fatal("curated entry should win over the registry duplicate") + } + } + if results := SearchMarket(catalog, "pager", 5); len(results) == 0 || results[0].Name != "pager-mcp" { + t.Fatalf("search = %+v", results) + } + if results := SearchMarket(catalog, "", 3); len(results) != 3 || results[0].Name != "datadog" { + t.Fatalf("empty search = %+v", results) + } + if results := SearchMarket(catalog, "zzzzqqq", 5); len(results) != 0 { + t.Fatalf("no-match search = %+v", results) + } +} + +func TestIndexCacheRoundTrip(t *testing.T) { + path := filepath.Join(t.TempDir(), "mcp-index.json") + now := time.Now().UTC().Truncate(time.Second) + if err := SaveIndex(path, []MarketEntry{{Name: "x", Config: uvxConfig("x")}}, now); err != nil { + t.Fatal(err) + } + entries, indexedAt, err := LoadIndex(path) + if err != nil || len(entries) != 1 || !indexedAt.Equal(now) || entries[0].Config.Command != "uvx" { + t.Fatalf("entries=%v at=%v err=%v", entries, indexedAt, err) + } +} + +func TestAddFromMarketWritesUserScopeSwitchedOff(t *testing.T) { + home := t.TempDir() + cwd := t.TempDir() + store := &Store{UserPath: filepath.Join(home, "mcp.json"), CWD: cwd} + states := &StateStore{Path: filepath.Join(home, "mcp-state.json")} + entry := MarketEntry{Name: "pager-mcp", Config: npxConfig("@acme/pager-mcp")} + server, err := AddFromMarket(store, states, entry) + if err != nil { + t.Fatal(err) + } + users, _ := store.Servers(ScopeUser) + if len(users) != 1 || users[0].Name != "pager-mcp" || users[0].ID() != server.ID() { + t.Fatalf("user servers = %+v", users) + } + if project, _ := store.Servers(ScopeProject); len(project) != 0 { + t.Fatalf("marketplace add leaked into project scope: %+v", project) + } + state, _ := states.Load() + if state.IsEnabled(users[0]) { + t.Fatal("marketplace server must stay off until the user turns it on") + } + if _, err := AddFromMarket(store, states, entry); err == nil { + t.Fatal("duplicate name should fail") + } +} + +func TestPartialRegistryFetchIsNotCached(t *testing.T) { + server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + if r.URL.Query().Get("cursor") == "c2" { + w.WriteHeader(http.StatusInternalServerError) + return + } + _, _ = w.Write([]byte(registryPage1)) + })) + defer server.Close() + entries, err := FetchRegistry(context.Background(), server.Client(), server.URL+"/v0/servers", nil) + if err == nil || len(entries) == 0 { + t.Fatalf("partial fetch = %d entries, err %v; want entries and an error", len(entries), err) + } + indexPath := filepath.Join(t.TempDir(), "mcp-index.json") + catalog, err := LoadCatalog(context.Background(), server.Client(), server.URL+"/v0/servers", indexPath, true, nil) + if err == nil || len(catalog) == 0 { + t.Fatalf("LoadCatalog = %d entries, err %v; want partial results and an error", len(catalog), err) + } + if _, statErr := os.Stat(indexPath); !os.IsNotExist(statErr) { + t.Fatalf("partial index was cached: %v", statErr) + } +} + +func TestManagerAddFromMarketRejectsImportedNames(t *testing.T) { + home := t.TempDir() + project := t.TempDir() + if err := os.WriteFile(filepath.Join(home, ".claude.json"), []byte(`{"mcpServers":{"pager-mcp":{"command":"pager"}}}`), 0o600); err != nil { + t.Fatal(err) + } + store := &Store{UserPath: filepath.Join(home, "mcp.json"), CWD: project} + states := &StateStore{Path: filepath.Join(home, "mcp-state.json")} + manager := NewManager(ManagerOptions{CWD: project, Store: store, State: states, Discovery: &Discovery{Home: home, CWD: project}}) + defer manager.Close() + if _, err := manager.AddFromMarket(MarketEntry{Name: "pager-mcp", Config: npxConfig("@acme/pager-mcp")}); err == nil { + t.Fatal("a marketplace entry must not reuse an imported server's name") + } + if users, _ := store.Servers(ScopeUser); len(users) != 0 { + t.Fatalf("user servers = %+v", users) + } +} + +func TestAddCommandIsShellQuoted(t *testing.T) { + remote := MarketEntry{Name: "acme", Config: ServerConfig{Type: "http", URL: "https://example.com/mcp?a=1&b=2"}} + if got := AddCommand(remote); got != "rootly mcp add -s user -t http acme 'https://example.com/mcp?a=1&b=2'" { + t.Fatalf("remote = %s", got) + } + local := MarketEntry{Name: "fs", Config: ServerConfig{Command: "npx", Args: []string{"-y", "@acme/fs", "${PWD}", "it's"}}} + if got := AddCommand(local); got != `rootly mcp add -s user fs -- npx -y @acme/fs '${PWD}' 'it'\''s'` { + t.Fatalf("local = %s", got) + } +} + +func TestCatalogKeepsRegistryEntriesWithTheSameShortName(t *testing.T) { + catalog := Catalog([]MarketEntry{ + {Name: "docs", RegistryName: "io.github.alpha/docs", Config: npxConfig("@alpha/docs")}, + {Name: "docs", RegistryName: "io.github.beta/docs", Config: npxConfig("@beta/docs")}, + {Name: "docs", Config: npxConfig("@gamma/docs")}, + }) + var names []string + for _, entry := range catalog { + if strings.HasPrefix(entry.Name, "docs") { + names = append(names, entry.Name) + } + } + if strings.Join(names, ",") != "docs,docs-beta,docs-2" { + t.Fatalf("names = %v", names) + } + for _, name := range names { + if err := ValidateServerName(name); err != nil { + t.Fatalf("%s: %v", name, err) + } + } +} + +func TestCatalogOnlyLetsCuratedReplaceTheSameServer(t *testing.T) { + catalog := Catalog([]MarketEntry{ + {Name: "linear", RegistryName: "app.linear/linear", Config: ServerConfig{Type: TransportHTTP, URL: "https://mcp.linear.app/mcp"}}, + {Name: "linear", RegistryName: "io.github.someone/linear", Config: npxConfig("@someone/linear-mcp")}, + }) + var linear []MarketEntry + for _, entry := range catalog { + if strings.HasPrefix(entry.Name, "linear") { + linear = append(linear, entry) + } + } + if len(linear) != 2 || linear[0].Name != "linear" || linear[0].RegistryName != "" || linear[1].Name != "linear-someone" { + t.Fatalf("linear entries = %+v, want curated linear plus linear-someone", linear) + } +} diff --git a/internal/mcp/resources.go b/internal/mcp/resources.go new file mode 100644 index 0000000..33d040b --- /dev/null +++ b/internal/mcp/resources.go @@ -0,0 +1,148 @@ +package mcp + +import ( + "context" + "errors" + "fmt" + "strings" + "time" + "unicode/utf8" + + sdk "github.com/modelcontextprotocol/go-sdk/mcp" +) + +const ( + maxDiscovered = 500 + // ResourceByteCap bounds how much of one resource is attached to a prompt. + ResourceByteCap = 32 * 1024 +) + +type PromptArg struct { + Name string + Required bool +} + +type PromptInfo struct { + Name string + Description string + Arguments []PromptArg +} + +type ResourceInfo struct { + URI string + Name string + Description string + MIMEType string +} + +// discoverExtras lists prompts and resources, only for the capabilities the server declares. +func discoverExtras(ctx context.Context, session *sdk.ClientSession) (prompts []PromptInfo, resources []ResourceInfo) { + init := session.InitializeResult() + if init == nil || init.Capabilities == nil { + return nil, nil + } + if init.Capabilities.Prompts != nil { + for prompt, err := range session.Prompts(ctx, nil) { + if err != nil || len(prompts) >= maxDiscovered { + break + } + info := PromptInfo{Name: prompt.Name, Description: prompt.Description} + for _, argument := range prompt.Arguments { + info.Arguments = append(info.Arguments, PromptArg{Name: argument.Name, Required: argument.Required}) + } + prompts = append(prompts, info) + } + } + if init.Capabilities.Resources != nil { + for resource, err := range session.Resources(ctx, nil) { + if err != nil || len(resources) >= maxDiscovered { + break + } + resources = append(resources, ResourceInfo{URI: resource.URI, Name: resource.Name, Description: resource.Description, MIMEType: resource.MIMEType}) + } + } + return prompts, resources +} + +func (m *Manager) sessionFor(serverID string) (*sdk.ClientSession, time.Duration, error) { + m.mu.Lock() + defer m.mu.Unlock() + for _, runtime := range m.servers { + if runtime.server.ID() == serverID && runtime.session != nil && !runtime.sessionOff { + return runtime.session, runtime.server.Config.CallTimeout(connectTimeout), nil + } + } + return nil, 0, fmt.Errorf("MCP server %s is not connected", serverID[strings.LastIndex(serverID, ":")+1:]) +} + +// GetPrompt renders an MCP prompt's text messages into one message. +func (m *Manager) GetPrompt(ctx context.Context, serverID, name string, arguments map[string]string) (string, error) { + session, timeout, err := m.sessionFor(serverID) + if err != nil { + return "", err + } + ctx, cancel := context.WithTimeout(ctx, timeout) + defer cancel() + result, err := session.GetPrompt(ctx, &sdk.GetPromptParams{Name: name, Arguments: arguments}) + if err != nil { + return "", err + } + var parts []string + for _, message := range result.Messages { + switch content := message.Content.(type) { + case *sdk.TextContent: + parts = append(parts, content.Text) + case *sdk.EmbeddedResource: + if content.Resource != nil && content.Resource.Text != "" { + parts = append(parts, content.Resource.Text) + } + } + } + text := strings.TrimSpace(strings.Join(parts, "\n\n")) + if text == "" { + return "", errors.New("prompt has no text content") + } + return text, nil +} + +// ReadResource returns a resource's text, cut to limit bytes on a rune boundary. +func (m *Manager) ReadResource(ctx context.Context, serverID, uri string, limit int) (text string, truncated bool, err error) { + session, timeout, err := m.sessionFor(serverID) + if err != nil { + return "", false, err + } + ctx, cancel := context.WithTimeout(ctx, timeout) + defer cancel() + result, err := session.ReadResource(ctx, &sdk.ReadResourceParams{URI: uri}) + if err != nil { + return "", false, err + } + var joined strings.Builder + for _, content := range result.Contents { + if content == nil || content.Text == "" { + continue + } + if joined.Len() > 0 { + joined.WriteByte('\n') + } + if limit > 0 && joined.Len()+len(content.Text) > limit { + joined.WriteString(content.Text[:min(len(content.Text), limit+1-joined.Len())]) + break + } + joined.WriteString(content.Text) + } + if joined.Len() == 0 { + return "", false, errors.New("no text content") + } + return capBytes(joined.String(), limit) +} + +func capBytes(text string, limit int) (capped string, truncated bool, err error) { + if limit <= 0 || len(text) <= limit { + return text, false, nil + } + for limit > 0 && !utf8.RuneStart(text[limit]) { + limit-- + } + return text[:limit], true, nil +} diff --git a/internal/mcp/resources_test.go b/internal/mcp/resources_test.go new file mode 100644 index 0000000..dc9e7dd --- /dev/null +++ b/internal/mcp/resources_test.go @@ -0,0 +1,150 @@ +package mcp + +import ( + "context" + "strings" + "testing" + + sdk "github.com/modelcontextprotocol/go-sdk/mcp" +) + +func connectInMemory(t *testing.T, server *sdk.Server) *sdk.ClientSession { + t.Helper() + ctx := context.Background() + clientTransport, serverTransport := sdk.NewInMemoryTransports() + serverSession, err := server.Connect(ctx, serverTransport, nil) + if err != nil { + t.Fatal(err) + } + t.Cleanup(func() { _ = serverSession.Close() }) + client := sdk.NewClient(&sdk.Implementation{Name: "test-client", Version: "1"}, nil) + session, err := client.Connect(ctx, clientTransport, nil) + if err != nil { + t.Fatal(err) + } + t.Cleanup(func() { _ = session.Close() }) + return session +} + +func docsServer() *sdk.Server { + server := sdk.NewServer(&sdk.Implementation{Name: "docs", Version: "1"}, nil) + server.AddPrompt(&sdk.Prompt{Name: "triage", Description: "Triage an alert", Arguments: []*sdk.PromptArgument{{Name: "service", Required: true}}}, + func(_ context.Context, request *sdk.GetPromptRequest) (*sdk.GetPromptResult, error) { + return &sdk.GetPromptResult{Messages: []*sdk.PromptMessage{{Role: "user", Content: &sdk.TextContent{Text: "Triage " + request.Params.Arguments["service"]}}}}, nil + }) + server.AddResource(&sdk.Resource{URI: "file:///runbook.md", Name: "runbook", MIMEType: "text/markdown"}, + func(_ context.Context, request *sdk.ReadResourceRequest) (*sdk.ReadResourceResult, error) { + return &sdk.ReadResourceResult{Contents: []*sdk.ResourceContents{{URI: request.Params.URI, Text: strings.Repeat("é", 40000)}}}, nil + }) + server.AddResource(&sdk.Resource{URI: "file:///logo.png", Name: "logo", MIMEType: "image/png"}, + func(_ context.Context, request *sdk.ReadResourceRequest) (*sdk.ReadResourceResult, error) { + return &sdk.ReadResourceResult{Contents: []*sdk.ResourceContents{{URI: request.Params.URI, Blob: []byte{1, 2, 3}}}}, nil + }) + return server +} + +func TestDiscoverExtrasListsPromptsAndResources(t *testing.T) { + session := connectInMemory(t, docsServer()) + prompts, resources := discoverExtras(context.Background(), session) + if len(prompts) != 1 || prompts[0].Name != "triage" || len(prompts[0].Arguments) != 1 || !prompts[0].Arguments[0].Required { + t.Fatalf("prompts = %+v", prompts) + } + if len(resources) != 2 || resources[0].URI != "file:///logo.png" && resources[1].URI != "file:///logo.png" { + t.Fatalf("resources = %+v", resources) + } +} + +func TestDiscoverExtrasSkipsUndeclaredCapabilities(t *testing.T) { + server := sdk.NewServer(&sdk.Implementation{Name: "tools-only", Version: "1"}, nil) + session := connectInMemory(t, server) + if caps := session.InitializeResult().Capabilities; caps.Prompts != nil || caps.Resources != nil { + t.Skip("SDK advertises prompts/resources without features") + } + prompts, resources := discoverExtras(context.Background(), session) + if prompts != nil || resources != nil { + t.Fatalf("prompts=%v resources=%v, want none", prompts, resources) + } +} + +func TestGetPromptAndReadResourceWithCap(t *testing.T) { + session := connectInMemory(t, docsServer()) + manager := &Manager{servers: []*runtimeServer{ + {server: Server{Name: "off"}, session: session, sessionOff: true}, + {server: Server{Name: "docs"}, session: session, status: StatusConnected}, + }} + ctx := context.Background() + docsID := Server{Name: "docs"}.ID() + text, err := manager.GetPrompt(ctx, docsID, "triage", map[string]string{"service": "api"}) + if err != nil || text != "Triage api" { + t.Fatalf("GetPrompt = %q, %v", text, err) + } + body, truncated, err := manager.ReadResource(ctx, docsID, "file:///runbook.md", 1001) + if err != nil || !truncated || len(body) != 1000 || !strings.HasPrefix(body, "éé") { + t.Fatalf("ReadResource = len %d truncated %v err %v", len(body), truncated, err) + } + full, truncated, err := manager.ReadResource(ctx, docsID, "file:///runbook.md", 0) + if err != nil || truncated || len(full) != 80000 { + t.Fatalf("uncapped ReadResource = len %d truncated %v err %v", len(full), truncated, err) + } + if _, _, err := manager.ReadResource(ctx, docsID, "file:///logo.png", ResourceByteCap); err == nil { + t.Fatal("binary resource should be unreadable") + } + if _, _, err := manager.ReadResource(ctx, "off", "file:///runbook.md", ResourceByteCap); err == nil { + t.Fatal("session-off server should not be readable") + } + if _, err := manager.GetPrompt(ctx, "missing", "triage", nil); err == nil { + t.Fatal("unknown server should fail") + } +} + +func TestSameNamedServersRouteByServerID(t *testing.T) { + other := connectInMemory(t, sdk.NewServer(&sdk.Implementation{Name: "other"}, nil)) + docs := connectInMemory(t, docsServer()) + imported := Server{Name: "docs", Scope: ScopeImported, Path: "/repo/.cursor/mcp.json"} + user := Server{Name: "docs", Scope: ScopeUser, Path: "/home/u/.rootly-cli/mcp.json"} + manager := &Manager{servers: []*runtimeServer{ + {server: imported, session: other, status: StatusConnected}, + {server: user, session: docs, status: StatusConnected}, + }} + ctx := context.Background() + if text, err := manager.GetPrompt(ctx, user.ID(), "triage", map[string]string{"service": "api"}); err != nil || text != "Triage api" { + t.Fatalf("GetPrompt = %q, %v", text, err) + } + if _, err := manager.GetPrompt(ctx, imported.ID(), "triage", nil); err == nil { + t.Fatal("GetPrompt by the imported server's ID must not reach the user server") + } + if _, _, err := manager.ReadResource(ctx, user.ID(), "file:///runbook.md", ResourceByteCap); err != nil { + t.Fatalf("ReadResource routed to the wrong server: %v", err) + } +} + +func TestLostAuthAfterConnectMarksNeedsAuth(t *testing.T) { + session := connectInMemory(t, docsServer()) + runtime := &runtimeServer{server: Server{Name: "remote"}, session: session, status: StatusConnected, generation: 3, tools: []*sdk.Tool{{Name: "x"}}} + manager := &Manager{servers: []*runtimeServer{runtime}, changes: make(chan struct{}, 1)} + probe := &authProbe{} + probe.watch(func() { manager.loseAuth(runtime, 3) }) + probe.mark() + if runtime.status != StatusNeedsAuth || runtime.session != nil || runtime.tools != nil { + t.Fatalf("runtime after 401 = status %s session %v tools %d", runtime.status, runtime.session, len(runtime.tools)) + } + select { + case <-manager.changes: + default: + t.Fatal("TUI was not notified") + } +} + +func TestAuthLostDuringDiscoveryMarksNeedsAuth(t *testing.T) { + session := connectInMemory(t, docsServer()) + runtime := &runtimeServer{server: Server{Name: "remote"}, session: session, status: StatusConnected, generation: 2} + manager := &Manager{servers: []*runtimeServer{runtime}, changes: make(chan struct{}, 1)} + probe := &authProbe{} + probe.mark() + manager.mu.Lock() + manager.watchAuth(runtime, 2, probe) + manager.mu.Unlock() + if runtime.status != StatusNeedsAuth || runtime.session != nil { + t.Fatalf("runtime after a 401 during discovery = status %s session %v", runtime.status, runtime.session) + } +} diff --git a/internal/mcp/security_test.go b/internal/mcp/security_test.go new file mode 100644 index 0000000..11d3847 --- /dev/null +++ b/internal/mcp/security_test.go @@ -0,0 +1,233 @@ +package mcp + +import ( + "context" + "encoding/json" + "net/http" + "net/http/httptest" + "path/filepath" + "strings" + "testing" +) + +func TestOriginMasksSecretArgs(t *testing.T) { + origin := ServerConfig{Command: "gh-mcp", Args: []string{"--token=sk-live", "--api-key", "abc", "serve"}}.Origin() + if strings.Contains(origin, "sk-live") || strings.Contains(origin, "abc") || !strings.Contains(origin, "serve") { + t.Fatalf("origin leaks secrets: %q", origin) + } +} + +func TestMaskToolArgumentsRedactsCommandAndCLISecrets(t *testing.T) { + values := map[string]any{ + "command": "export AWS_SECRET_ACCESS_KEY=aws-secret && gh --token 'cli-token'", + "args": []any{"--api-key", "api-secret", "AWS_ACCESS_KEY_ID=access-secret"}, + } + encoded, err := json.Marshal(MaskToolArguments(values)) + if err != nil { + t.Fatal(err) + } + for _, secret := range []string{"aws-secret", "cli-token", "api-secret", "access-secret"} { + if strings.Contains(string(encoded), secret) { + t.Errorf("masked arguments leaked %q: %s", secret, encoded) + } + } +} + +func TestMaskCommandMasksShellCodeInSecretValues(t *testing.T) { + for _, test := range []struct { + command string + code string + }{ + {command: `export TOKEN="$(rm x)"`, code: "rm x"}, + {command: "export TOKEN=\"`rm x`\"", code: "rm x"}, + {command: `export TOKEN="${rm x}"`, code: "rm x"}, + {command: `export TOKEN="<(rm x)"`, code: "rm x"}, + {command: `export TOKEN=">(rm x)"`, code: "rm x"}, + } { + masked := MaskCommand(test.command) + if strings.Contains(masked, test.code) || !strings.Contains(masked, Mask) { + t.Errorf("config command did not mask executable code: %q", masked) + } + } +} + +func TestMaskApprovalCommandKeepsExecutableCodeVisible(t *testing.T) { + for _, command := range []string{ + `export TOKEN="$(rm -rf ./data)"`, + "export TOKEN=\"`rm x`\"", + `export TOKEN="${rm x}"`, + `export TOKEN="<(rm x)"`, + `export TOKEN=">(rm x)"`, + } { + if got := MaskApprovalCommand(command); !strings.Contains(got, "rm ") { + t.Errorf("approval command hid executable code: %q", got) + } + } +} + +func TestMaskingConfigArgumentsMasksShellSyntax(t *testing.T) { + args := MaskCommandArgs([]string{"--password", "a$(b"}) + if strings.Contains(strings.Join(args, " "), "a$(b") || args[1] != Mask { + t.Fatalf("argv secret was not masked: %#v", args) + } + toolArguments := MaskToolArguments(map[string]any{"password": "x$(y)"}) + if toolArguments["password"] != Mask { + t.Fatalf("tool argument secret was not masked: %#v", toolArguments) + } + command := MaskCommand(`export TOKEN="$(rm -rf ./data)"`) + if strings.Contains(command, "rm -rf ./data") || !strings.Contains(command, Mask) { + t.Fatalf("command description did not mask executable code: %q", command) + } +} + +func TestMaskApprovalArgumentsPreservesShellCodeAtAnyDepth(t *testing.T) { + arguments := MaskApprovalArguments(map[string]any{ + "password": "x${y}", + "args": []any{"--password", "a$(b)", "TOKEN=$(rm x)"}, + "nested": map[string]any{"token": ">(rm x)"}, + }) + if arguments["password"] != "x${y}" { + t.Fatalf("approval arguments masked executable code in secret-named value: %#v", arguments) + } + args := arguments["args"].([]string) + if args[1] != "a$(b)" || args[2] != "TOKEN=$(rm x)" { + t.Fatalf("approval arguments masked executable code in argv: %#v", args) + } + nested := arguments["nested"].(map[string]any) + if nested["token"] != ">(rm x)" { + t.Fatalf("approval arguments masked executable code in nested value: %#v", nested) + } +} + +func TestServerHTTPClientKeepsCredentialsOnOrigin(t *testing.T) { + var leaked []string + other := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + leaked = append(leaked, r.Header.Get("X-Api-Key")+"|"+r.Header.Get("Authorization")) + w.WriteHeader(http.StatusOK) + })) + defer other.Close() + var seen string + server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + if r.URL.Path == "/redirect" { + http.Redirect(w, r, other.URL+"/steal", http.StatusFound) + return + } + seen = r.Header.Get("X-Api-Key") + "|" + r.Header.Get("Authorization") + w.WriteHeader(http.StatusOK) + })) + defer server.Close() + + client := newServerHTTPClient(ServerConfig{Type: TransportHTTP, URL: server.URL + "/mcp", Headers: map[string]string{"X-Api-Key": "k"}}, func(context.Context) string { return "tok" }, nil) + get := func(target string) error { + request, _ := http.NewRequestWithContext(t.Context(), http.MethodGet, target, http.NoBody) + response, err := client.Do(request) + if err == nil { + _ = response.Body.Close() + } + return err + } + if err := get(server.URL + "/mcp"); err != nil || seen != "k|Bearer tok" { + t.Fatalf("same origin should get credentials, err=%v seen=%q", err, seen) + } + if err := get(server.URL + "/redirect"); err == nil || !strings.Contains(err.Error(), "cross-origin redirect") { + t.Fatalf("cross-origin redirect should be refused, err=%v", err) + } + if err := get(other.URL + "/direct"); err != nil { + t.Fatal(err) + } + if len(leaked) != 1 || leaked[0] != "|" { + t.Fatalf("credentials reached another origin: %v", leaked) + } +} + +func TestImportedProjectFilesNeedApprovalPerHash(t *testing.T) { + home, project := t.TempDir(), t.TempDir() + writeFile(t, filepath.Join(home, ".claude.json"), marshalTestJSON(t, map[string]any{ + "mcpServers": map[string]any{ + "global": map[string]any{"command": "g"}, + }, + "projects": map[string]any{ + project: map[string]any{ + "mcpServers": map[string]any{ + "repo-db": map[string]any{"command": "db"}, + }, + }, + }, + })) + writeFile(t, filepath.Join(project, ".cursor", "mcp.json"), `{"mcpServers":{"cursor-tool":{"command":"run-me"}}}`) + writeFile(t, filepath.Join(project, ".vscode", "mcp.json"), `{"servers":{"vs-tool":{"command":"vs"}}}`) + store := &StateStore{Path: filepath.Join(home, "state.json")} + + servers, problems := Discovery{Home: home, CWD: project}.Imported() + if len(problems) != 0 { + t.Fatal(problems) + } + state, _ := store.Load() + byName := map[string]Server{} + for _, server := range servers { + byName[server.Name] = server + } + if state.NeedsApproval(byName["global"]) { + t.Fatal("user-level Claude server is not project-controlled") + } + for _, name := range []string{"repo-db", "cursor-tool", "vs-tool"} { + if !state.NeedsApproval(byName[name]) { + t.Fatalf("%s comes from the project and must need approval", name) + } + } + if err := store.Update(func(state *State) { state.Approve(byName["cursor-tool"]) }); err != nil { + t.Fatal(err) + } + state, _ = store.Load() + if state.NeedsApproval(byName["cursor-tool"]) { + t.Fatal("approved server should start") + } + + writeFile(t, filepath.Join(project, ".cursor", "mcp.json"), `{"mcpServers":{"cursor-tool":{"command":"evil"}}}`) + servers, _ = Discovery{Home: home, CWD: project}.Imported() + for _, server := range servers { + if server.Name == "cursor-tool" && !state.NeedsApproval(server) { + t.Fatal("a changed command must be approved again") + } + } +} + +func TestOriginOfNormalizesDefaultPorts(t *testing.T) { + if originOf("https://MCP.example.com:443/a") != originOf("https://mcp.example.com/b") { + t.Fatal("explicit default port should be the same origin") + } + if originOf("http://h:80") != originOf("http://h") || originOf("https://h:8443") == originOf("https://h") { + t.Fatal("port normalization is wrong") + } +} + +func TestLoadDoesNotReuseGlobalTrustForProjectServer(t *testing.T) { + home, project := t.TempDir(), t.TempDir() + claude := filepath.Join(home, ".claude.json") + writeFile(t, claude, `{"mcpServers":{"foo":{"command":"foo-not-installed"}}}`) + store := &StateStore{Path: filepath.Join(home, ".rootly-cli", "mcp-state.json")} + manager := NewManager(ManagerOptions{CWD: project, Store: &Store{UserPath: filepath.Join(home, ".rootly-cli", "mcp.json"), CWD: project}, State: store, Discovery: &Discovery{Home: home, CWD: project}}) + defer manager.Close() + _ = manager.Load() + if err := manager.Enable(manager.Views()[0].Server.ID()); err != nil { + t.Fatal(err) + } + if status := manager.Views()[0].Status; status == StatusPending || status == StatusOff { + t.Fatalf("enabled global server should start, got %s", status) + } + + writeFile(t, claude, marshalTestJSON(t, map[string]any{ + "projects": map[string]any{ + project: map[string]any{ + "mcpServers": map[string]any{ + "foo": map[string]any{"command": "foo-not-installed"}, + }, + }, + }, + })) + _ = manager.Load() + view := manager.Views()[0] + if view.Server.Project == "" || view.Status != StatusPending { + t.Fatalf("same-named project server must need approval, got project=%q status=%s", view.Server.Project, view.Status) + } +} diff --git a/internal/mcp/sources.go b/internal/mcp/sources.go new file mode 100644 index 0000000..40add4f --- /dev/null +++ b/internal/mcp/sources.go @@ -0,0 +1,181 @@ +package mcp + +import ( + "encoding/json" + "os" + "path/filepath" + "runtime" + "sort" + + toml "github.com/pelletier/go-toml/v2" +) + +// ImportSource is a third-party MCP config rootly can read but never writes. +type ImportSource struct { + Label string + Path string + parse func([]byte, string) (map[string]ServerConfig, error) + // project marks a file inside the checkout; projectNames picks the + // project-scoped entries of a file that mixes both. + project bool + projectNames func([]byte, string) map[string]bool +} + +type Discovery struct { + Home string + CWD string + GOOS string +} + +func NewDiscovery(cwd string) Discovery { + home, _ := os.UserHomeDir() + return Discovery{Home: home, CWD: cwd, GOOS: runtime.GOOS} +} + +func (d Discovery) Sources() []ImportSource { + sources := []ImportSource{ + {Label: "Claude Code", Path: filepath.Join(d.Home, ".claude.json"), parse: parseClaudeCode, projectNames: claudeProjectNames}, + {Label: "Codex", Path: filepath.Join(d.Home, ".codex", "config.toml"), parse: parseCodexTOML}, + {Label: "Cursor", Path: filepath.Join(d.Home, ".cursor", "mcp.json"), parse: parseMCPServersKey}, + {Label: "Cursor (project)", Path: filepath.Join(d.CWD, ".cursor", "mcp.json"), parse: parseMCPServersKey, project: true}, + {Label: "VS Code (project)", Path: filepath.Join(d.CWD, ".vscode", "mcp.json"), parse: parseVSCode, project: true}, + } + switch d.GOOS { + case "darwin": + sources = append(sources, ImportSource{Label: "Claude Desktop", Path: filepath.Join(d.Home, "Library", "Application Support", "Claude", "claude_desktop_config.json"), parse: parseMCPServersKey}) + case "linux": + sources = append(sources, ImportSource{Label: "Claude Desktop", Path: filepath.Join(d.Home, ".config", "Claude", "claude_desktop_config.json"), parse: parseMCPServersKey}) + } + return sources +} + +type ImportError struct { + Source ImportSource + Err error +} + +// Imported lists servers from every readable third-party config. Missing +// files are skipped; malformed files are reported without failing the rest. +func (d Discovery) Imported() ([]Server, []ImportError) { + var servers []Server + var problems []ImportError + for _, source := range d.Sources() { + data, err := os.ReadFile(source.Path) + if err != nil { + if !os.IsNotExist(err) { + problems = append(problems, ImportError{Source: source, Err: err}) + } + continue + } + parsed, err := source.parse(data, d.CWD) + if err != nil { + problems = append(problems, ImportError{Source: source, Err: err}) + continue + } + listed := sortedServers(parsed, ScopeImported, source.Path, source.Label) + var fromProject map[string]bool + if source.projectNames != nil { + fromProject = source.projectNames(data, d.CWD) + } + for index := range listed { + if source.project || fromProject[listed[index].Name] { + listed[index].Project = d.CWD + } + } + servers = append(servers, listed...) + } + return servers, problems +} + +func parseMCPServersKey(data []byte, _ string) (map[string]ServerConfig, error) { + document := jsonDocument{} + if err := json.Unmarshal(data, &document); err != nil { + return nil, err + } + return decodeServersLenient(document["mcpServers"]), nil +} + +func parseClaudeCode(data []byte, cwd string) (map[string]ServerConfig, error) { + var document struct { + MCPServers json.RawMessage `json:"mcpServers"` + Projects map[string]map[string]json.RawMessage `json:"projects"` + } + if err := json.Unmarshal(data, &document); err != nil { + return nil, err + } + servers := decodeServersLenient(document.MCPServers) + for name, server := range decodeServersLenient(document.Projects[cwd]["mcpServers"]) { + servers[name] = server + } + return servers, nil +} + +func claudeProjectNames(data []byte, cwd string) map[string]bool { + var document struct { + Projects map[string]map[string]json.RawMessage `json:"projects"` + } + if err := json.Unmarshal(data, &document); err != nil { + return nil + } + names := map[string]bool{} + for name := range decodeServersLenient(document.Projects[cwd]["mcpServers"]) { + names[name] = true + } + return names +} + +func parseVSCode(data []byte, _ string) (map[string]ServerConfig, error) { + document := jsonDocument{} + if err := json.Unmarshal(data, &document); err != nil { + return nil, err + } + return decodeServersLenient(document["servers"]), nil +} + +type codexServer struct { + Command string `toml:"command"` + Args []string `toml:"args"` + Env map[string]string `toml:"env"` + URL string `toml:"url"` + BearerTokenEnvVar string `toml:"bearer_token_env_var"` + HTTPHeaders map[string]string `toml:"http_headers"` + EnvHTTPHeaders map[string]string `toml:"env_http_headers"` + Enabled *bool `toml:"enabled"` +} + +func parseCodexTOML(data []byte, _ string) (map[string]ServerConfig, error) { + var document struct { + MCPServers map[string]codexServer `toml:"mcp_servers"` + } + if err := toml.Unmarshal(data, &document); err != nil { + return nil, err + } + servers := make(map[string]ServerConfig, len(document.MCPServers)) + names := make([]string, 0, len(document.MCPServers)) + for name := range document.MCPServers { + names = append(names, name) + } + sort.Strings(names) + for _, name := range names { + entry := document.MCPServers[name] + if entry.Enabled != nil && !*entry.Enabled { + continue + } + server := ServerConfig{Command: entry.Command, Args: entry.Args, Env: entry.Env} + if entry.URL != "" { + server = ServerConfig{Type: TransportHTTP, URL: entry.URL, BearerTokenEnvVar: entry.BearerTokenEnvVar} + headers := map[string]string{} + for key, value := range entry.HTTPHeaders { + headers[key] = value + } + for key, variable := range entry.EnvHTTPHeaders { + headers[key] = "${" + variable + "}" + } + if len(headers) > 0 { + server.Headers = headers + } + } + servers[name] = server + } + return servers, nil +} diff --git a/internal/mcp/state.go b/internal/mcp/state.go new file mode 100644 index 0000000..e4c01e9 --- /dev/null +++ b/internal/mcp/state.go @@ -0,0 +1,154 @@ +package mcp + +import ( + "encoding/json" + "os" + "path/filepath" + "sync" + + "github.com/rootlyhq/rootly-cli/internal/config" +) + +const stateFile = "mcp-state.json" + +type Approval struct { + Project string `json:"project"` + Server string `json:"server"` + Hash string `json:"hash"` +} + +// State is rootly-owned MCP state: which imported servers are enabled and +// which project .mcp.json servers were approved. +type State struct { + EnabledImported []string `json:"enabled_imported,omitempty"` + Disabled []string `json:"disabled,omitempty"` + Approvals []Approval `json:"approvals,omitempty"` + BuiltinTools map[string]bool `json:"builtin_tools,omitempty"` +} + +type StateStore struct { + Path string + mu sync.Mutex +} + +func NewStateStore() *StateStore { + return &StateStore{Path: filepath.Join(config.Dir(), stateFile)} +} + +func (s *StateStore) Load() (State, error) { + s.mu.Lock() + defer s.mu.Unlock() + return s.load() +} + +func (s *StateStore) load() (State, error) { + var state State + data, err := os.ReadFile(s.Path) + if err != nil { + if os.IsNotExist(err) { + return state, nil + } + return state, err + } + if len(data) == 0 { + return state, nil + } + err = json.Unmarshal(data, &state) + return state, err +} + +func (s *StateStore) Update(change func(*State)) error { + s.mu.Lock() + defer s.mu.Unlock() + state, err := s.load() + if err != nil { + return err + } + change(&state) + document := jsonDocument{} + encoded, err := json.Marshal(state) + if err != nil { + return err + } + if err := json.Unmarshal(encoded, &document); err != nil { + return err + } + return writeJSONDocument(s.Path, document, 0o600) +} + +func (state State) IsApproved(server Server) bool { + for _, approval := range state.Approvals { + if approval.Project == server.approvalPath() && approval.Server == server.Name && approval.Hash == server.Config.Hash() { + return true + } + } + return false +} + +func (state *State) Approve(server Server) { + kept := state.Approvals[:0] + for _, approval := range state.Approvals { + if approval.Project != server.approvalPath() || approval.Server != server.Name { + kept = append(kept, approval) + } + } + kept = append(kept, Approval{Project: server.approvalPath(), Server: server.Name, Hash: server.Config.Hash()}) + state.Approvals = kept +} + +func (state State) IsEnabled(server Server) bool { + if server.Scope == ScopeImported { + return containsString(state.EnabledImported, server.ID()) + } + return !containsString(state.Disabled, server.ID()) +} + +func (state *State) SetEnabled(server Server, enabled bool) { + if server.Scope == ScopeImported { + state.EnabledImported = setString(state.EnabledImported, server.ID(), enabled) + return + } + state.Disabled = setString(state.Disabled, server.ID(), !enabled) +} + +func (state State) IsBuiltinEnabled(name string) bool { + if enabled, ok := state.BuiltinTools[name]; ok { + return enabled + } + return name != "local_bash" +} + +func (state *State) SetBuiltinEnabled(name string, enabled bool) { + if state.BuiltinTools == nil { + state.BuiltinTools = make(map[string]bool) + } + state.BuiltinTools[name] = enabled +} + +func containsString(values []string, value string) bool { + for _, candidate := range values { + if candidate == value { + return true + } + } + return false +} + +func setString(values []string, value string, present bool) []string { + kept := make([]string, 0, len(values)+1) + for _, candidate := range values { + if candidate != value { + kept = append(kept, candidate) + } + } + if present { + kept = append(kept, value) + } + return kept +} + +// NeedsApproval reports whether a project-controlled server still needs the +// user's approval of its current config before it may start. +func (state State) NeedsApproval(server Server) bool { + return server.NeedsProjectTrust() && !state.IsApproved(server) +}