1package mcp23import (4 "bytes"5 "context"6 "encoding/json"7 "errors"8 "io"9 "os"10 "path/filepath"11 "slices"12 "strings"13 "testing"1415 sdk "github.com/modelcontextprotocol/go-sdk/mcp"1617 "src.rest/cat/012/internal/headless"18 "src.rest/cat/012/internal/notebook"19 "src.rest/cat/012/internal/nushell"20 "src.rest/cat/012/internal/sheet"21)2223// bookFile writes a workbook: a sheet of sales with a header, and a24// second sheet whose name has a space.25func bookFile(t *testing.T) string {26 t.Helper()27 w := sheet.NewBook()28 s := w.Sheet(0)29 rows := [][]string{{"Region", "Units", "Price"}, {"North", "3", "12.5"}, {"South", "5", "9.75"}, {"North", "7", "12.5"}}30 for r, row := range rows {31 for c, in := range row {32 s.Set(sheet.Addr{Col: c, Row: r}, in)33 }34 }35 other, _ := w.AddSheet("Q3 plan", 1)36 other.Set(sheet.Addr{}, "=Sheet1!B2*2")37 path := filepath.Join(t.TempDir(), "book.012")38 writeBook(t, w, path)39 return path40}4142func writeBook(t *testing.T, w *sheet.Workbook, path string) {43 t.Helper()44 var b bytes.Buffer45 if err := w.Write(&b); err != nil {46 t.Fatal(err)47 }48 if err := os.WriteFile(path, b.Bytes(), 0o644); err != nil {49 t.Fatal(err)50 }51}5253// connect serves o to a client over the SDK's in-memory transport. A54// server with a default workbook and no roots is open to its folder.55func connect(t *testing.T, o Options) *sdk.ClientSession {56 t.Helper()57 return connectClient(t, o, sdk.NewClient(&sdk.Implementation{Name: "test", Version: "1"}, nil), nil)58}5960// connectClient serves o to client c, connecting with co.61func connectClient(t *testing.T, o Options, c *sdk.Client, co *sdk.ClientSessionOptions) *sdk.ClientSession {62 t.Helper()63 if o.Roots == nil && o.Default != "" {64 roots, err := NewRoots([]string{filepath.Dir(o.Default)})65 if err != nil {66 t.Fatal(err)67 }68 o.Roots = roots69 }70 ctx := context.Background()71 st, ct := sdk.NewInMemoryTransports()72 ss, err := New(o).Connect(ctx, st, nil)73 if err != nil {74 t.Fatal(err)75 }76 t.Cleanup(func() { ss.Close() })77 cs, err := c.Connect(ctx, ct, co)78 if err != nil {79 t.Fatal(err)80 }81 t.Cleanup(func() { cs.Close() })82 return cs83}8485// call calls a tool, failing on a protocol error, and decodes its86// structured result into out.87func call(t *testing.T, cs *sdk.ClientSession, name string, args any, out any) *sdk.CallToolResult {88 t.Helper()89 res, err := cs.CallTool(context.Background(), &sdk.CallToolParams{Name: name, Arguments: args})90 if err != nil {91 t.Fatalf("%s: %v", name, err)92 }93 if out != nil && !res.IsError {94 data, _ := json.Marshal(res.StructuredContent)95 if err := json.Unmarshal(data, out); err != nil {96 t.Fatalf("%s: %v in %s", name, err, data)97 }98 }99 return res100}101102func errorText(res *sdk.CallToolResult) string {103 if !res.IsError || len(res.Content) == 0 {104 return ""105 }106 return res.Content[0].(*sdk.TextContent).Text107}108109func TestTools(t *testing.T) {110 cs := connect(t, Options{Default: bookFile(t)})111 var names []string112 for tl, err := range cs.Tools(context.Background(), nil) {113 if err != nil {114 t.Fatal(err)115 }116 names = append(names, tl.Name)117 }118 for _, want := range []string{"describe", "read_range", "evaluate", "find", "list_errors", "write_cells", "apply_operations",119 "sort", "filter", "create_chart", "create_pivot"} {120 if !slices.Contains(names, want) {121 t.Errorf("no tool %s in %v", want, names)122 }123 }124 if slices.Contains(names, "run_notebook_cell") {125 t.Error("run_notebook_cell offered without Notebooks")126 }127 ro := connect(t, Options{Default: bookFile(t), ReadOnly: true})128 for tl := range ro.Tools(context.Background(), nil) {129 if !tl.Annotations.ReadOnlyHint {130 t.Errorf("read-only server offers %s", tl.Name)131 }132 }133}134135func TestDescribeAndRead(t *testing.T) {136 cs := connect(t, Options{Default: bookFile(t)})137 var d headless.Description138 call(t, cs, "describe", map[string]any{}, &d)139 if len(d.Sheets) != 2 || d.Sheets[0].Header != 1 || !slices.Equal(d.Sheets[0].Columns, []string{"Region", "Units", "Price"}) {140 t.Errorf("describe: %+v", d.Sheets)141 }142 var r headless.Range143 call(t, cs, "read_range", map[string]any{"ref": "A1:C4", "text": true}, &r)144 if r.Range != "Sheet1!A1:C4" || r.Rows != 4 || string(r.Values[1][2]) != "12.5" || r.Text[0][0] != "Region" {145 t.Errorf("read_range: %+v", r)146 }147 call(t, cs, "read_range", map[string]any{"ref": "'Q3 plan'!A1"}, &r)148 if string(r.Values[0][0]) != "6" || r.Formulas["A1"] != "=Sheet1!B2*2" {149 t.Errorf("read_range of a formula: %+v", r)150 }151 res := call(t, cs, "read_range", map[string]any{"ref": "Q4!A1"}, nil)152 if !strings.Contains(errorText(res), `no sheet named "Q4" (sheets: Sheet1, Q3 plan)`) {153 t.Errorf("bad sheet: %q", errorText(res))154 }155 var found headless.Found156 call(t, cs, "find", map[string]any{"query": "north"}, &found)157 if found.Total != 2 || found.Matches[1].Cell != "Sheet1!A4" {158 t.Errorf("find: %+v", found)159 }160}161162func TestEvaluateLeavesTheFile(t *testing.T) {163 path := bookFile(t)164 before, _ := os.ReadFile(path)165 cs := connect(t, Options{Default: path})166 var ev headless.Evaluated167 call(t, cs, "evaluate", map[string]any{"formula": `SUMIF(A2:A4, "North", B2:B4)`}, &ev)168 if string(ev.Value) != "10" || ev.Cell != "Sheet1!A6" {169 t.Errorf("evaluate: %+v", ev)170 }171 call(t, cs, "evaluate", map[string]any{"formula": "=SEQUENCE(2)", "at": "E1"}, &ev)172 if ev.Spill == nil || ev.Spill.Range != "Sheet1!E1:E2" {173 t.Errorf("spill: %+v", ev.Spill)174 }175 call(t, cs, "evaluate", map[string]any{"formula": "=1/0"}, &ev)176 if string(ev.Value) != `"#DIV/0!"` || ev.Error == "" {177 t.Errorf("error: %+v", ev)178 }179 if res := call(t, cs, "evaluate", map[string]any{"formula": "=SUM(A1"}, nil); !strings.Contains(errorText(res), "Expected") {180 t.Errorf("parse error: %q", errorText(res))181 }182 var errs errorsOut183 call(t, cs, "list_errors", map[string]any{}, &errs)184 if len(errs.Errors) != 0 {185 t.Errorf("errors: %+v", errs)186 }187 if after, _ := os.ReadFile(path); !bytes.Equal(before, after) {188 t.Error("reading tools changed the file")189 }190}191192func TestWrites(t *testing.T) {193 path := bookFile(t)194 cs := connect(t, Options{Default: path})195 before, _ := os.ReadFile(path)196 var ch Changed197 call(t, cs, "write_cells", map[string]any{"entries": []map[string]string{{"ref": "D1", "input": "Total"}, {"ref": "D2", "input": "=B2*C2"}},198 "dry_run": true}, &ch)199 if ch.Saved || len(ch.Changes) != 3 || ch.Changes[1].Item != "D2" || string(ch.Changes[2].New) != "37.5" {200 t.Errorf("dry run: %+v", ch)201 }202 if after, _ := os.ReadFile(path); !bytes.Equal(before, after) {203 t.Error("dry run saved")204 }205 call(t, cs, "write_cells", map[string]any{"entries": []map[string]string{{"ref": "D2", "input": "=B2*C2"}}}, &ch)206 if !ch.Saved || len(ch.Changes) != 2 {207 t.Errorf("write: %+v", ch)208 }209 f, _ := headless.Open(path, false)210 if v := f.Book.Sheet(0).Value(sheet.Addr{Col: 3, Row: 1}); v.Num != 37.5 {211 t.Errorf("D2 = %v", v)212 }213 res := call(t, cs, "write_cells", map[string]any{"entries": []map[string]string{{"ref": "D3", "input": "=SUM(A1"}}}, nil)214 if !strings.Contains(errorText(res), "Sheet1!D3: Expected") {215 t.Errorf("bad formula: %q", errorText(res))216 }217 call(t, cs, "apply_operations", map[string]any{"operations": []map[string]any{218 {"op": "insert_rows", "ref": "A2"}, {"op": "set", "ref": "A2", "input": "East"}, {"op": "define_name", "ref": "B2:B5", "name": "Units"},219 }}, &ch)220 if !ch.Saved {221 t.Errorf("operations: %+v", ch)222 }223 call(t, cs, "sort", map[string]any{"ref": "A1:C5", "header": true, "keys": []map[string]any{{"column": "Units", "descending": true}}}, &ch)224 var r headless.Range225 call(t, cs, "read_range", map[string]any{"ref": "Units"}, &r)226 if got := string(r.Values[0][0]) + string(r.Values[1][0]); got != "75" {227 t.Errorf("sorted units start %s: %+v", got, r.Values)228 }229 call(t, cs, "filter", map[string]any{"ref": "A1:C5", "columns": []map[string]any{{"column": "Region", "condition": "eq", "value": "North"}}}, &ch)230 var d headless.Description231 call(t, cs, "describe", map[string]any{}, &d)232 if d.Sheets[0].Filter != "A1:C5" || len(d.Names) != 1 {233 t.Errorf("filter or name missing: %+v", d)234 }235}236237func TestChartAndPivot(t *testing.T) {238 path := bookFile(t)239 cs := connect(t, Options{Default: path})240 var c chartOut241 call(t, cs, "create_chart", map[string]any{"data": "A1:B4", "type": "bar", "title": "Units"}, &c)242 if c.Chart.Number != 1 || c.Chart.Sheet != "Sheet1" || !c.Saved {243 t.Errorf("chart: %+v", c)244 }245 var p pivotOut246 call(t, cs, "create_pivot", map[string]any{"source": "A1:C4", "rows": []string{"Region"}, "values": []map[string]any{{"column": "Units"}}}, &p)247 if p.Sheet != "Pivot Table 1" || !p.Saved {248 t.Errorf("pivot: %+v", p)249 }250 var r headless.Range251 call(t, cs, "read_range", map[string]any{"ref": "'Pivot Table 1'!A1:B3"}, &r)252 if string(r.Values[1][0]) != `"North"` || string(r.Values[1][1]) != "10" {253 t.Errorf("pivot result: %+v", r.Values)254 }255 res := call(t, cs, "create_pivot", map[string]any{"source": "A1:C4", "values": []map[string]any{{"column": "Weight"}}}, nil)256 if !strings.Contains(errorText(res), `"Weight" isn't a column of A1:C4`) {257 t.Errorf("unknown column: %q", errorText(res))258 }259}260261func TestResourcesAndPrompts(t *testing.T) {262 cs := connect(t, Options{Default: bookFile(t)})263 ctx := context.Background()264 var uris []string265 for r, err := range cs.Resources(ctx, nil) {266 if err != nil {267 t.Fatal(err)268 }269 uris = append(uris, r.URI)270 }271 if !slices.Contains(uris, "o12://book.012/Sheet1!A1:C4") || !slices.Contains(uris, "o12://book.012/Q3%20plan!A1") {272 t.Errorf("resources: %v", uris)273 }274 for uri, want := range map[string]string{275 "o12://book.012/Sheet1!A1:C4": `"range":"Sheet1!A1:C4"`,276 "o12://book.012/Q3%20plan!A1": `"formulas":{"A1":"=Sheet1!B2*2"}`,277 "o12://book.012/Sheet1!B2:B3": `"values":[[3],[5]]`,278 "o12://book.012/table/Missing": "isn't a cell",279 } {280 res, err := cs.ReadResource(ctx, &sdk.ReadResourceParams{URI: uri})281 got := ""282 if err != nil {283 got = err.Error()284 } else {285 got = res.Contents[0].Text286 }287 if !strings.Contains(got, want) {288 t.Errorf("%s: %s, want %s", uri, got, want)289 }290 }291 var names []string292 for p, err := range cs.Prompts(ctx, nil) {293 if err != nil {294 t.Fatal(err)295 }296 names = append(names, p.Name)297 }298 if len(names) != len(prompts) {299 t.Errorf("prompts: %v", names)300 }301 res, err := cs.GetPrompt(ctx, &sdk.GetPromptParams{Name: "add_column", Arguments: map[string]string{"range": "Sheet1!A1:C4", "column": "revenue"}})302 if err != nil {303 t.Fatal(err)304 }305 if text := res.Messages[0].Content.(*sdk.TextContent).Text; !strings.Contains(text, "Add a column to Sheet1!A1:C4 in the workbook (describe without a path lists them) that computes revenue") {306 t.Errorf("prompt: %s", text)307 }308}309310var errNotTrusted = errors.New("book.012 was saved on another computer")311312// fakeNu answers every cell with one table.313type fakeNu struct{ ran int }314315func (f *fakeNu) Run(_ context.Context, _ nushell.Job, _ string, stdout io.Writer) error {316 f.ran++317 _, err := stdout.Write([]byte("[[name]; [a.txt]]"))318 return err319}320321func TestRunNotebookCell(t *testing.T) {322 w := sheet.NewBook()323 nb, _ := w.AddNotebook("Notes", w.Sheet(0))324 nb.SetNotebookCells("add cell", []notebook.Cell{{Kind: notebook.Note, Source: "# Files"}, {Kind: notebook.Code, Source: "files = ls"}})325 path := filepath.Join(t.TempDir(), "nb.012")326 writeBook(t, w, path)327 nu := &fakeNu{}328 refuse := true329 cs := connect(t, Options{Default: path, Notebooks: &headless.NotebookOptions{Runner: nu},330 MayRun: func(string, *sheet.Workbook) error {331 if refuse {332 return errNotTrusted333 }334 return nil335 }})336 if res := call(t, cs, "run_notebook_cell", map[string]any{"notebook": "Notes", "cell": 2}, nil); errorText(res) != errNotTrusted.Error() || nu.ran != 0 {337 t.Errorf("untrusted: %q, ran %d", errorText(res), nu.ran)338 }339 refuse = false340 var run headless.CellRun341 call(t, cs, "run_notebook_cell", map[string]any{"notebook": "Notes", "cell": 2}, &run)342 if run.State != "ran" || run.Name != "files" || !strings.Contains(run.Output, "a.txt") {343 t.Errorf("run: %+v", run)344 }345 res, err := cs.ReadResource(context.Background(), &sdk.ReadResourceParams{URI: "o12://nb.012/notebook/Notes/2"})346 if err != nil || !strings.Contains(res.Contents[0].Text, "a.txt") {347 t.Errorf("cell resource after the run: %v %+v", err, res)348 }349 if res := call(t, cs, "run_notebook_cell", map[string]any{"notebook": "Notes", "cell": 1}, nil); !strings.Contains(errorText(res), "is a note") {350 t.Errorf("note cell: %q", errorText(res))351 }352}