1package ui23import (4 "context"5 "errors"6 "fmt"7 "log/slog"8 "slices"9 "time"1011 tea "charm.land/bubbletea/v2"1213 "src.rest/cat/012/internal/notebook"14 "src.rest/cat/012/internal/nushell"15 "src.rest/cat/012/internal/sheet"16 "src.rest/cat/012/internal/ui/nbview"17)1819// Running cells. One runs at a time, in the background, so the screen20// stays live and Stop kills it: the cells asked for are a queue, each21// started when the one before has finished. What a cell reads (other22// cells' outputs, ranges of sheets) is taken as it starts; what it23// printed becomes its output as it ends, and goes on to the sheet it was24// sent to. A failure stops the queue, as Jupyter's run all stops.2526// nbState is the notebooks' part of the model.27type nbState struct {28 views map[*sheet.Sheet]*nbview.View29 runner nushell.Runner // nu, or a fake in tests30 // runs are the cells running and waiting: the model's own, or in a31 // shared room the room's (share.go).32 runs *nbRuns33 clip []notebook.Cell34 // words are nu's command names, asked for once (asked).35 words []string36 asked bool37 // served is set in a session of 012 serve, where cells run only38 // when the server's config allows them.39 served bool40 // startup shows the notebook once the program starts (012 nu).41 startup bool42 // fromFile is set when the workbook opened with cells of its own,43 // which may have been made on another computer.44 fromFile bool45 // grid is the sheet shown last that isn't a notebook, and sel its46 // selection then: what $selection reads.47 grid *sheet.Sheet48 sel sheet.Rect49 // stale caches which cells are stale, for the state it was worked50 // out in.51 stale map[int]bool52 staleKey [2]int53 // still stops the running cell's head turning, so tests that run54 // commands in line don't wait on its ticks.55 still bool56 // saved is the outputs' change count when the file was last saved,57 // and saving when it was being saved.58 saved, saving int59 // lang asks nu about the cells being written (nblang.go), and60 // langKey is what the words it was last given were worked out from.61 lang *nbview.NuSession62 langKey nbLangKey63 // follow scrolls the view to each cell as it starts, while a run of64 // several cells goes on and the user hasn't scrolled away.65 follow bool66 // streamTick is set while streaming cells' heads are ticking:67 // nbstream.go.68 streamTick bool69 // out is the outputs' grids (nbgrid.go), and frame counts the frames70 // drawn, which the grids clear what they keep for a frame by.71 out outGrids72 frame int73}7475// nbRuns are the cells running, waiting and streaming. In a room of76// 012 serve they are the room's, which every participant sees and any77// may stop, and whose results reach whoever keeps the room then78// (room.Seat.Post), so they outlast the participant who started them.79type nbRuns struct {80 queue []nbQueued81 running *nbRun82 gen int83 count int // runs of this workbook, for [n]84 // streams are the cells running as streams, by cell ID: nbstream.go.85 streams map[int]*nbStream86}8788// nbQueued is a cell waiting to run.89type nbQueued struct {90 s *sheet.Sheet91 id int92}9394// nbRun is a cell running.95type nbRun struct {96 nbQueued97 gen int98 cancel context.CancelFunc99 start time.Time100 source string101 reads map[string]int102 selection string103}104105type (106 // nbDoneMsg is a run finished: what it printed, or why it failed.107 nbDoneMsg struct {108 run *nbRun109 nuon []byte110 vars map[string][]byte // the variables it handed back111 err error112 }113 // nbWordsMsg brings nu's command names.114 nbWordsMsg struct{ words []string }115 // nbTickMsg redraws a running cell's head.116 nbTickMsg struct{ gen int }117)118119// SetShellRunner makes cells run, and the code editor ask, with r120// rather than nu, for tests; it's set before anything is asked.121func (m *Model) SetShellRunner(r nushell.Runner) {122 m.nb.runner = r123 if m.nb.lang != nil {124 m.nb.lang.Runner = r125 }126}127128func (m *Model) runner() nushell.Runner {129 if m.nb.runner == nil {130 return nushell.Nu{}131 }132 return m.nb.runner133}134135// cellState is how cell id of notebook s stands.136func (m *Model) cellState(s *sheet.Sheet, id int) nbview.State {137 var st nbview.State138 if r := m.nb.runs.running; r != nil && r.s == s && r.id == id {139 st.Running, st.Started = true, r.start140 }141 if r := m.nb.runs.streams[id]; r != nil && r.s == s {142 st.Running, st.Started, st.Live, st.Rows = true, r.start, true, r.rows143 }144 st.Waiting = slices.Contains(m.nb.runs.queue, nbQueued{s, id})145 st.Stale = m.staleCells(s)[id]146 cells := s.NotebookCells()147 if i := slices.IndexFunc(cells, func(c notebook.Cell) bool { return c.ID == id }); i >= 0 {148 st.Problem = notebook.Taken(cells, i)149 }150 return st151}152153// staleCells are notebook s's stale cells, worked out once a state.154func (m *Model) staleCells(s *sheet.Sheet) map[int]bool {155 key := [2]int{m.book().StateID(), m.book().OutputsChanged()}156 if m.nb.stale == nil || m.nb.staleKey != key {157 m.nb.stale = map[int]bool{}158 for _, t := range m.book().Sheets() {159 if t.IsNotebook() {160 for id := range notebook.Stale(t.NotebookCells(), m.book().Output) {161 m.nb.stale[id] = true162 }163 }164 }165 m.nb.staleKey = key166 }167 return m.nb.stale168}169170// runSelected runs the selected cells: then, with next 1, selects the171// cell below them (adding one at the end), or with 2 adds one under172// them, as Jupyter's Shift+Enter and Alt+Enter.173func (m *Model) runSelected(next int) tea.Cmd {174 v := m.nbView()175 v.StopEdit()176 from, to := v.Range()177 cells := m.sheet.NotebookCells()178 if from >= len(cells) {179 return nil180 }181 var cmd tea.Cmd182 switch {183 case from < to:184 cmd = m.runRange(from, to+1)185 case cells[from].Kind == notebook.Code:186 cmd = m.runCells(m.sheet, from, cells)187 }188 switch {189 case next == 2 || next == 1 && to == len(cells)-1:190 return tea.Batch(cmd, m.addCell(1, next == 2))191 case next == 1:192 v.Select(to+1, false)193 }194 return cmd195}196197// runRange runs the code cells from index from up to to (-1 for the198// last), each after the cells it reads, the view following the cell199// running.200func (m *Model) runRange(from, to int) tea.Cmd {201 v := m.nbView()202 v.StopEdit()203 cells := m.sheet.NotebookCells()204 if to < 0 {205 to = len(cells)206 }207 var want []int208 for i := from; i < to && i < len(cells); i++ {209 want = append(want, i)210 }211 m.nb.follow = len(want) > 1212 return m.queueCells(m.sheet, cells, want)213}214215// runCells runs cell i, after what it reads that hasn't run, and in a216// reactive notebook then the cells reading it.217func (m *Model) runCells(s *sheet.Sheet, i int, cells []notebook.Cell) tea.Cmd {218 want := append(notebook.Inputs(cells, i, m.book().Output), i)219 if s.Reactive() {220 want = append(want, notebook.Dependents(cells, i)...)221 }222 return m.queueCells(s, cells, want)223}224225// queueCells queues the code cells of want, in the order they run, asking226// first when the cells came from another computer.227func (m *Model) queueCells(s *sheet.Sheet, cells []notebook.Cell, want []int) tea.Cmd {228 order, cycle := notebook.Order(cells, want)229 for _, i := range cycle {230 m.fail(fmt.Sprintf("Cell %d reads itself, through the cells it reads: it doesn't run", i+1))231 }232 if len(order) == 0 {233 return nil234 }235 return m.trustNotebook(func(m *Model) tea.Cmd {236 for _, i := range order {237 q := nbQueued{s, cells[i].ID}238 if !slices.Contains(m.nb.runs.queue, q) && (m.nb.runs.running == nil || m.nb.runs.running.nbQueued != q) {239 m.nb.runs.queue = append(m.nb.runs.queue, q)240 }241 }242 if m.nb.runs.running != nil {243 return nil244 }245 return m.nextCell()246 })247}248249// shellOff says why cells can't run here, or "".250func (m *Model) shellOff() string {251 switch {252 case m.nb.served && !m.configBool("serve-shell", false):253 return "Notebook cells don't run in 012 serve (serve-shell in the server's config)"254 case m.configString("shell", "ask") == "off":255 return "Notebook cells don't run (shell = off in your config)"256 }257 return ""258}259260// trustNotebook runs run once cells may run: shell = on, or the file's261// cells were made or trusted here, or the user agrees now.262func (m *Model) trustNotebook(run func(*Model) tea.Cmd) tea.Cmd {263 if why := m.shellOff(); why != "" {264 m.fail(why)265 return nil266 }267 if m.configString("shell", "ask") == "on" || m.macroTrusted() {268 return run(m)269 }270 m.ask(question{271 msg: "Run this file's notebook cells?",272 desc: "It was saved on another computer: its cells run as you, with your files",273 choices: []choice{274 {key: "enter", label: "Run", run: func(m *Model) tea.Cmd { m.trustHere(); return run(m) }},275 {key: "esc", label: "Cancel", run: func(*Model) tea.Cmd { return nil }},276 },277 })278 return nil279}280281// nextCell starts the first cell queued.282func (m *Model) nextCell() tea.Cmd {283 for len(m.nb.runs.queue) > 0 {284 q := m.nb.runs.queue[0]285 m.nb.runs.queue = m.nb.runs.queue[1:]286 c, ok := cellOf(q)287 if !ok || c.Kind != notebook.Code {288 continue289 }290 cells := q.s.NotebookCells()291 i := slices.IndexFunc(cells, func(x notebook.Cell) bool { return x.ID == c.ID })292 if why := notebook.Taken(cells, i); why != "" {293 m.failCell(q, c, why)294 continue295 }296 m.stopStream(q.id) // run once, it no longer follows297 run := &nbRun{nbQueued: q, source: c.Source, start: time.Now()}298 job, err := m.jobFor(q.s, c, run, false)299 if err != nil {300 m.failCell(q, c, err.Error())301 continue302 }303 return m.startCell(run, job)304 }305 m.nb.follow = false // nothing left to follow306 return nil307}308309// cellOf is the cell q names, if it's still there.310func cellOf(q nbQueued) (notebook.Cell, bool) {311 for _, c := range q.s.NotebookCells() {312 if c.ID == q.id {313 return c, true314 }315 }316 return notebook.Cell{}, false317}318319// failCell gives a cell that can't run an output saying why, and stops320// the queue.321func (m *Model) failCell(q nbQueued, c notebook.Cell, why string) {322 m.nb.runs.count++323 m.book().SetOutput(q.id, ¬ebook.Output{Err: why, Count: m.nb.runs.count, Source: c.Source})324 m.nb.runs.queue = nil325}326327// startCell runs a cell in the background.328func (m *Model) startCell(run *nbRun, job nushell.Job) tea.Cmd {329 m.nb.runs.gen++330 ctx, cancel := context.WithCancel(context.Background())331 run.gen, run.cancel = m.nb.runs.gen, cancel332 m.nb.runs.running = run333 if m.nb.follow {334 m.viewOf(run.s).Reveal(run.id) // Run all: the view follows the cell running335 }336 runner, timeout, parent := m.runner(), m.configDuration("nu-timeout", 30*time.Second), m.spans.Parent()337 exec := func() tea.Msg {338 defer cancel()339 span := parent.Start("nu", slog.Int("tables", len(job.Tables)))340 out, vars, err := nushell.ExecVars(ctx, runner, job, timeout, maxOutput)341 if err != nil {342 span.Fail(err)343 } else {344 span.End(slog.Int("bytes", len(out)))345 }346 return nbDoneMsg{run: run, nuon: out, vars: vars, err: err}347 }348 return tea.Batch(m.roomOwned(exec), m.tickCell(run.gen))349}350351// maxOutput is the most a run may print: past it the run is stopped.352const maxOutput = 512 << 20353354// tickCell redraws the running cell's head every tenth of a second.355func (m *Model) tickCell(gen int) tea.Cmd {356 if m.nb.still {357 return nil358 }359 return tea.Tick(100*time.Millisecond, func(time.Time) tea.Msg { return nbTickMsg{gen} })360}361362// ticked keeps the running cell's head turning.363func (m *Model) ticked(msg nbTickMsg) tea.Cmd {364 if r := m.nb.runs.running; r != nil && r.gen == msg.gen {365 return m.tickCell(msg.gen)366 }367 return nil368}369370// finishCell keeps what a run made, sends it on to its sheet, and371// starts the next cell queued.372func (m *Model) finishCell(msg nbDoneMsg) tea.Cmd {373 r := msg.run374 if m.nb.runs.running != r {375 return nil // stopped, and another started since376 }377 m.nb.runs.running = nil378 c, ok := cellOf(r.nbQueued)379 if !ok {380 return m.nextCell()381 }382 m.nb.runs.count++383 o := ¬ebook.Output{Count: m.nb.runs.count, Took: time.Since(r.start), Source: r.source, Reads: r.reads, Selection: r.selection}384 if msg.err != nil {385 o.Err = runError(msg.err)386 var nuErr *nushell.Error387 if errors.As(msg.err, &nuErr) {388 o.Detail = nuErr.Help()389 }390 m.nb.runs.queue = nil391 } else {392 o.NUON, o.Vars = msg.nuon, msg.vars393 if o.NUON == nil {394 o.NUON = []byte{}395 }396 }397 m.book().SetOutput(c.ID, o)398 m.feedOutput(c.Name())399 if msg.err == nil && r.s.Reactive() {400 cells := r.s.NotebookCells()401 if i := slices.IndexFunc(cells, func(x notebook.Cell) bool { return x.ID == c.ID }); i >= 0 {402 for _, j := range notebook.Dependents(cells, i) {403 if q := (nbQueued{r.s, cells[j].ID}); !slices.Contains(m.nb.runs.queue, q) {404 m.nb.runs.queue = append(m.nb.runs.queue, q)405 }406 }407 }408 }409 if len(m.nb.runs.queue) == 0 && time.Since(r.start) > 5*time.Second {410 return tea.Batch(m.nextCell(), m.term.notify("Notebook cells finished in "+m.displayName()))411 }412 return m.nextCell()413}414415// runError is what a failed run says.416func runError(err error) string {417 var nuErr *nushell.Error418 switch {419 case errors.Is(err, context.Canceled):420 return "stopped"421 case errors.As(err, &nuErr):422 return nuErr.Msg423 }424 return err.Error()425}426427// stopCells stops the cell running, killing its process, forgets428// those waiting, and stops every stream.429func (m *Model) stopCells() {430 if r := m.nb.runs.running; r != nil {431 r.cancel()432 }433 m.nb.runs.queue = nil434 m.stopStreams()435}436437// clearOutputs clears every output of the notebook shown; restart also438// stops what's running and counts runs from 1 again.439func (m *Model) clearOutputs(restart bool) {440 if restart {441 m.stopCells()442 m.nb.runs.running = nil443 m.nb.runs.count = 0444 }445 for _, c := range m.sheet.NotebookCells() {446 if m.book().Output(c.ID) != nil {447 m.book().SetOutput(c.ID, nil)448 m.feedOutput(c.Name())449 }450 }451 m.note = "Cleared every output"452 if restart {453 m.note = "Restarted: every output cleared, runs counted from 1"454 }455}