1package cowork23import (4 "context"5 "errors"6 "strings"7 "testing"8 "time"910 sdk "github.com/modelcontextprotocol/go-sdk/mcp"1112 "src.rest/cat/012/internal/headless"13 "src.rest/cat/012/internal/mcp"14 "src.rest/cat/012/internal/notebook"15 "src.rest/cat/012/internal/room"16 "src.rest/cat/012/internal/sheet"17)1819// person is a participant standing for a session.20type person struct{ name string }2122func (p *person) Name() string { return p.name }23func (p *person) Notify() {}2425// bot is one of the agents participant without a connection.26type bot struct{ person }2728func (*bot) Agent() {}2930// liveRoom is a room with ann in it, on a workbook of 1, 2, 3 in A1:A331// and their sum in A4, and the agent attached.32func liveRoom(t *testing.T) (ann *room.Seat, agent *Agent) {33 t.Helper()34 g := room.NewRegistry(room.Edit)35 ann, _ = g.Join("k", &person{"ann"}, func() *sheet.Workbook {36 w := sheet.New().Book()37 for i, in := range []string{"1", "2", "3", "=SUM(A1:A3)"} {38 w.Sheet(0).Set(sheet.Addr{Row: i}, in)39 }40 w.AddSheet("Q3", 1)41 return w42 })43 ann.Do(func(*sheet.Workbook) { ann.SetSaved(room.Saved{Name: "budget.012"}) })44 seat, ok := g.JoinOpen("k", &bot{person{"claude"}})45 if !ok {46 t.Fatal("the agent couldn't join")47 }48 return ann, NewAgent(seat)49}5051func set(entries ...string) func(*sheet.Workbook) error {52 return func(w *sheet.Workbook) error {53 var es []headless.Entry54 for i := 0; i+1 < len(entries); i += 2 {55 es = append(es, headless.Entry{Ref: entries[i], Input: entries[i+1]})56 }57 _, err := headless.Set(w, es, headless.SetOptions{})58 return err59 }60}6162func value(s *room.Seat, sheetName, ref string) sheet.Value {63 var v sheet.Value64 s.Do(func(w *sheet.Workbook) {65 a, _ := sheet.ParseAddr(ref)66 v = w.Lookup(sheetName).Value(a)67 })68 return v69}7071func TestSuggestionWaitsForThePerson(t *testing.T) {72 ann, a := liveRoom(t)73 made, err := a.Make(context.Background(), mcp.Making{Label: "set", Message: "fix the totals", Fn: set("A1", "10", "A2", "20")})74 if err != nil {75 t.Fatal(err)76 }77 if made.Applied || made.Suggestion != 1 || len(made.Changes) == 0 {78 t.Fatalf("made %+v", made)79 }80 if v := value(ann, "Sheet1", "A4"); v.Num != 6 {81 t.Fatalf("the suggestion was made: A4 is %v", v)82 }83 ann.Do(func(*sheet.Workbook) {84 b := BoardOf(ann)85 s, i, ok := b.At("Sheet1", sheet.Addr{Row: 1})86 if !ok || s.ID != 1 || i != 1 || s.Message != "fix the totals" || s.Agent != "claude" {87 t.Errorf("A2's suggestion: %+v %d %v", s, i, ok)88 }89 })90}9192func TestSuggestionAcceptedIsTheAgentsStep(t *testing.T) {93 ann, a := liveRoom(t)94 a.Make(context.Background(), mcp.Making{Label: "set", Fn: set("A1", "10", "A2", "20")})95 ann.Do(func(w *sheet.Workbook) {96 if err := BoardOf(ann).Accept(w, 1, nil); err != nil {97 t.Fatal(err)98 }99 if w.CanUndo() {100 t.Error("ann can undo the agent's step")101 }102 })103 if v := value(ann, "Sheet1", "A4"); v.Num != 33 {104 t.Fatalf("A4 is %v after accepting", v)105 }106 news := a.report()107 if len(news) != 1 || !strings.Contains(news[0], "was accepted") {108 t.Fatalf("news %v", news)109 }110 if label, err := a.Undo(context.Background()); err != nil || label != "set" {111 t.Fatalf("the agent's undo: %q %v", label, err)112 }113 if v := value(ann, "Sheet1", "A4"); v.Num != 6 {114 t.Fatalf("A4 is %v after the agent's undo", v)115 }116}117118func TestSuggestionByCell(t *testing.T) {119 ann, a := liveRoom(t)120 a.Make(context.Background(), mcp.Making{Label: "set", Fn: set("A1", "10", "A2", "20", "A3", "30")})121 ann.Do(func(w *sheet.Workbook) {122 b := BoardOf(ann)123 if err := b.Accept(w, 1, []int{0}); err != nil {124 t.Fatal(err)125 }126 if err := b.Reject(1, []int{1}); err != nil {127 t.Fatal(err)128 }129 if !b.Get(1).Pending() {130 t.Error("A3 isn't pending")131 }132 b.Accept(w, 1, nil) // the rest: A3133 })134 if v := value(ann, "Sheet1", "A4"); v.Num != 42 {135 t.Fatalf("A4 is %v: 10 + 2 + 30", v)136 }137 list, _ := a.Suggestions(context.Background())138 if len(list) != 1 || list[0].State != "accepted in part" || list[0].Accepted != 2 || list[0].Rejected != 1 {139 t.Fatalf("suggestions %+v", list)140 }141}142143func TestSuggestionRejected(t *testing.T) {144 ann, a := liveRoom(t)145 a.Make(context.Background(), mcp.Making{Label: "set", Fn: set("A1", "10")})146 ann.Do(func(*sheet.Workbook) { BoardOf(ann).Reject(1, nil) })147 made, _ := a.Make(context.Background(), mcp.Making{Label: "set", Fn: set("B1", "x")})148 if len(made.Notices) != 1 || !strings.Contains(made.Notices[0], "Suggestion 1 (set) was rejected") {149 t.Fatalf("notices %v", made.Notices)150 }151 if v := value(ann, "Sheet1", "A1"); v.Num != 1 {152 t.Fatal("a rejected suggestion changed A1")153 }154 ann.Do(func(w *sheet.Workbook) {155 if err := BoardOf(ann).Accept(w, 1, nil); !errors.Is(err, ErrGone) {156 t.Errorf("accepting a rejected suggestion: %v", err)157 }158 })159}160161func TestDirectEdits(t *testing.T) {162 ann, a := liveRoom(t)163 ann.Do(func(*sheet.Workbook) { BoardOf(ann).Direct = true })164 made, err := a.Make(context.Background(), mcp.Making{Label: "set", Fn: set("A1", "10")})165 if err != nil || !made.Applied || made.Suggestion != 0 {166 t.Fatalf("made %+v, %v", made, err)167 }168 if v := value(ann, "Sheet1", "A4"); v.Num != 15 {169 t.Fatalf("A4 is %v", v)170 }171 var where room.Presence172 ann.Do(func(*sheet.Workbook) { where = ann.Peers()[0].Presence })173 if where.Cursor != (sheet.Addr{}) || where.Sheet == nil {174 t.Errorf("the agent's pointer isn't on its change: %+v", where)175 }176}177178func TestScopeRefuses(t *testing.T) {179 for _, tc := range []struct {180 name string181 scope func(w *sheet.Workbook) Scope182 fn func(*sheet.Workbook) error183 want string184 }{185 {"read only", func(*sheet.Workbook) Scope { return Scope{Kind: ReadOnly} }, set("A1", "5"), "read only"},186 {"another sheet", func(w *sheet.Workbook) Scope { return Scope{Kind: OneSheet, Sheet: w.Lookup("Q3")} }, set("A1", "5"), "Sheet1!A1, outside what the person let you change (sheet Q3)"},187 {"outside the range", func(w *sheet.Workbook) Scope {188 return Scope{Kind: OneRange, Sheet: w.Sheet(0), Range: sheet.Rect{To: sheet.Addr{Row: 2}}}189 }, set("B1", "5"), "Sheet1!B1, outside what the person let you change (Sheet1!A1:A3)"},190 {"rows in a range", func(w *sheet.Workbook) Scope {191 return Scope{Kind: OneRange, Sheet: w.Sheet(0), Range: sheet.Rect{To: sheet.Addr{Row: 2}}}192 }, func(w *sheet.Workbook) error { return w.Sheet(0).InsertRows(0, 1) }, "inserts or deletes rows"},193 {"a sheet added", func(w *sheet.Workbook) Scope { return Scope{Kind: OneSheet, Sheet: w.Sheet(0)} },194 func(w *sheet.Workbook) error { _, err := w.AddSheet("New", 2); return err }, "adds, deletes, renames or moves sheets"},195 } {196 t.Run(tc.name, func(t *testing.T) {197 ann, a := liveRoom(t)198 ann.Do(func(w *sheet.Workbook) { BoardOf(ann).Scope = tc.scope(w) })199 _, err := a.Make(context.Background(), mcp.Making{Label: "x", Fn: tc.fn})200 if err == nil || !strings.Contains(err.Error(), tc.want) {201 t.Fatalf("got %v, want %q", err, tc.want)202 }203 })204 }205 ann, a := liveRoom(t)206 ann.Do(func(w *sheet.Workbook) { BoardOf(ann).Scope = Scope{Kind: OneSheet, Sheet: w.Sheet(0)} })207 if _, err := a.Make(context.Background(), mcp.Making{Label: "x", Fn: set("C3", "5")}); err != nil {208 t.Fatalf("inside the sheet: %v", err)209 }210 if _, err := a.Focus(context.Background(), "Q3!A1"); err == nil {211 t.Fatal("focus outside the scope")212 }213 if at, err := a.Focus(context.Background(), "B2:C3"); err != nil || at != "Sheet1!B2:C3" {214 t.Fatalf("focus: %q %v", at, err)215 }216}217218// answer answers the first question on the board as a session would.219func answer(t *testing.T, ann *room.Seat, ans Answer) *Ask {220 t.Helper()221 deadline := time.Now().Add(5 * time.Second)222 for time.Now().Before(deadline) {223 var q *Ask224 ann.Do(func(*sheet.Workbook) {225 b := BoardOf(ann)226 if q = b.NextAsk(ann.ID(), func(int) bool { return true }); q != nil {227 b.Answer(q, ans)228 }229 })230 if q != nil {231 return q232 }233 time.Sleep(5 * time.Millisecond)234 }235 t.Fatal("no question came")236 return nil237}238239func TestAskReturnsTheAnswer(t *testing.T) {240 ann, a := liveRoom(t)241 got := make(chan *sdk.ElicitResult, 1)242 go func() {243 res, err := a.Ask(context.Background(), &sdk.ElicitParams{Message: "Overwrite B2:B40?",244 RequestedSchema: map[string]any{"type": "object", "properties": map[string]any{245 "sure": map[string]any{"type": "boolean", "title": "Sure"},246 "which": map[string]any{"type": "string", "enum": []string{"a", "b"}}}, "required": []string{"which"}}})247 if err != nil {248 t.Error(err)249 }250 got <- res251 }()252 q := answer(t, ann, Answer{Action: "accept", Content: map[string]any{"sure": true, "which": "b"}})253 if q.Message != "Overwrite B2:B40?" || len(q.Fields) != 2 || q.Fields[0].Name != "which" || q.Fields[0].Kind != Choice || q.Fields[1].Kind != Boolean {254 t.Fatalf("question %+v", q)255 }256 res := <-got257 if res.Action != "accept" || res.Content["which"] != "b" {258 t.Fatalf("result %+v", res)259 }260}261262func TestAskWithdrawnWhenTheAgentGivesUp(t *testing.T) {263 ann, a := liveRoom(t)264 ctx, cancel := context.WithCancel(context.Background())265 done := make(chan error, 1)266 go func() {267 _, err := a.Ask(ctx, &sdk.ElicitParams{Message: "Sure?"})268 done <- err269 }()270 for n := 0; n == 0; {271 ann.Do(func(*sheet.Workbook) { n = len(BoardOf(ann).Asks()) })272 time.Sleep(time.Millisecond)273 }274 cancel()275 if err := <-done; !errors.Is(err, context.Canceled) {276 t.Fatal(err)277 }278 ann.Do(func(*sheet.Workbook) {279 if n := len(BoardOf(ann).Asks()); n != 0 {280 t.Errorf("%d questions left", n)281 }282 })283}284285func TestAskRefusesSchemasItCantShow(t *testing.T) {286 _, a := liveRoom(t)287 _, err := a.Ask(context.Background(), &sdk.ElicitParams{Message: "Where?",288 RequestedSchema: map[string]any{"type": "object", "properties": map[string]any{"at": map[string]any{"type": "array"}}}})289 if err == nil || !strings.Contains(err.Error(), "questions take strings") {290 t.Fatal(err)291 }292}293294func TestJEVDirectAsksFirst(t *testing.T) {295 ann, a := liveRoom(t)296 ann.Do(func(*sheet.Workbook) { BoardOf(ann).Direct = true })297 done := make(chan error, 1)298 go func() {299 _, err := a.Make(context.Background(), mcp.Making{Label: "ask", Fn: set("B1", `=JEV.TEST(A1, "odd?")`)})300 done <- err301 }()302 q := answer(t, ann, Answer{Action: "accept", Always: true})303 if q.Kind != Grant || q.Grant != "jev" {304 t.Fatalf("question %+v", q)305 }306 if err := <-done; err != nil {307 t.Fatal(err)308 }309 ann.Do(func(*sheet.Workbook) {310 if !BoardOf(ann).Grants.JEV {311 t.Error("always didn't grant JEV for the session")312 }313 })314}315316func TestNotebookRunsAskAndGoToTheKeeper(t *testing.T) {317 ann, a := liveRoom(t)318 ann.Do(func(w *sheet.Workbook) { w.AddNotebook("Notes", w.Sheet(0)) })319 ann.Do(func(w *sheet.Workbook) {320 w.Lookup("Notes").SetNotebookCells("add", []notebook.Cell{{Source: "ls"}})321 })322 if _, err := a.RunCell(context.Background(), "Nope", 1); err == nil {323 t.Fatal("ran a cell of no notebook")324 }325 done := make(chan error, 1)326 go func() {327 _, err := a.RunCell(context.Background(), "Notes", 1)328 done <- err329 }()330 q := answer(t, ann, Answer{Action: "decline"})331 if q.Grant != "notebooks" {332 t.Fatalf("question %+v", q)333 }334 if err := <-done; err == nil {335 t.Fatal("a declined run went ahead")336 }337}