1package cowork23import (4 "encoding/json"5 "errors"6 "fmt"7 "slices"8 "strconv"9 "strings"10)1112// Questions: the agent asks the person something with the ask tool,13// shaped as an MCP elicitation request (a message, and a flat object14// schema of what to answer), and the person answers on the context line15// of their session, one field at a time. Elicitation goes from a server16// to its host's user; in live mode the person is on the server's side,17// so the agent asks through a tool, and the answer comes back as an18// elicitation's result: accept with the values, decline or cancel.19// Grants (running notebook cells, JEV) are asked the same way.2021// AskKind is what a question is for.22type AskKind int2324const (25 // Question is the agent's own question (the ask tool).26 Question AskKind = iota27 // Grant asks leave to run notebook cells or enter JEV formulas.28 Grant29)3031// FieldKind is the type of what a field asks for.32type FieldKind int3334const (35 Text FieldKind = iota36 Number37 Integer38 Boolean39 Choice40)4142// Field is one thing a question asks for.43type Field struct {44 Name, Title, Description string45 Kind FieldKind46 Choices []string // a Choice's values47 Labels []string // and how they're shown, when the schema names them48 Required bool49}5051// Label is how the field is named to the person.52func (f Field) Label() string {53 if f.Title != "" {54 return f.Title55 }56 return f.Name57}5859// Ask is a question waiting for a person's answer.60type Ask struct {61 ID int62 Author int // the agent's seat63 Agent string64 Color int65 Kind AskKind66 Message string67 Fields []Field68 // Grant is what a Grant asks for: "notebooks" or "jev".69 Grant string7071 reply chan Answer72 taken int // the seat showing it, 0 for none yet73 done bool74}7576// Answer is the person's answer: the elicitation's action (accept,77// decline, cancel) and values by field, and for a grant whether it78// holds for the rest of the session.79type Answer struct {80 Action string81 Content map[string]any82 Always bool83}8485// NewAsk makes a question, its fields read from an elicitation's86// requested schema (nil for a yes or no question).87func NewAsk(message string, schema any) (*Ask, error) {88 if strings.TrimSpace(message) == "" {89 return nil, errors.New("ask needs a message: the question shown to the person")90 }91 fields, err := Fields(schema)92 if err != nil {93 return nil, err94 }95 return &Ask{Kind: Question, Message: message, Fields: fields, reply: make(chan Answer, 1)}, nil96}9798// NewGrant asks leave for what ("notebooks" or "jev").99func NewGrant(what, message string) *Ask {100 return &Ask{Kind: Grant, Grant: what, Message: message, reply: make(chan Answer, 1)}101}102103// Reply is where the answer arrives, once.104func (a *Ask) Reply() <-chan Answer { return a.reply }105106// Done reports whether the question was answered or withdrawn.107func (a *Ask) Done() bool { return a.done }108109// schemaProp is a property of an elicitation's schema, as far as110// questions read it.111type schemaProp struct {112 Type string `json:"type"`113 Title string `json:"title"`114 Description string `json:"description"`115 Enum []string `json:"enum"`116 EnumNames []string `json:"enumNames"`117 OneOf []struct {118 Const string `json:"const"`119 Title string `json:"title"`120 } `json:"oneOf"`121}122123// Fields reads an elicitation's requested schema: an object of124// top-level string, number, integer and boolean properties, strings with125// enum (or oneOf consts) being a choice, as MCP allows.126func Fields(schema any) ([]Field, error) {127 if schema == nil {128 return nil, nil129 }130 data, err := json.Marshal(schema)131 if err != nil {132 return nil, err133 }134 var s struct {135 Type string `json:"type"`136 Properties map[string]json.RawMessage `json:"properties"`137 Required []string `json:"required"`138 }139 if err := json.Unmarshal(data, &s); err != nil {140 return nil, fmt.Errorf("requestedSchema: %w", err)141 }142 if s.Type != "" && s.Type != "object" {143 return nil, errors.New("requestedSchema must be an object schema")144 }145 names := make([]string, 0, len(s.Properties))146 for n := range s.Properties {147 names = append(names, n)148 }149 slices.Sort(names) // JSON objects are unordered; required ones first below150 slices.SortStableFunc(names, func(a, b string) int {151 return boolRank(slices.Contains(s.Required, b)) - boolRank(slices.Contains(s.Required, a))152 })153 var out []Field154 for _, n := range names {155 f, err := field(n, s.Properties[n])156 if err != nil {157 return nil, err158 }159 f.Required = slices.Contains(s.Required, n)160 out = append(out, f)161 }162 return out, nil163}164165func boolRank(b bool) int {166 if b {167 return 1168 }169 return 0170}171172// field reads one property.173func field(name string, raw json.RawMessage) (Field, error) {174 var p schemaProp175 if err := json.Unmarshal(raw, &p); err != nil {176 return Field{}, fmt.Errorf("requestedSchema property %s: %w", name, err)177 }178 f := Field{Name: name, Title: p.Title, Description: p.Description}179 switch {180 case p.Type == "string" && len(p.Enum) > 0:181 f.Kind, f.Choices, f.Labels = Choice, p.Enum, p.EnumNames182 case p.Type == "string" && len(p.OneOf) > 0:183 f.Kind = Choice184 for _, o := range p.OneOf {185 f.Choices, f.Labels = append(f.Choices, o.Const), append(f.Labels, o.Title)186 }187 case p.Type == "string":188 f.Kind = Text189 case p.Type == "number":190 f.Kind = Number191 case p.Type == "integer":192 f.Kind = Integer193 case p.Type == "boolean":194 f.Kind = Boolean195 default:196 return Field{}, fmt.Errorf("requestedSchema property %s is a %q: questions take strings, numbers, integers, booleans and choices", name, p.Type)197 }198 if len(f.Choices) > 9 {199 return Field{}, fmt.Errorf("requestedSchema property %s has %d choices: at most 9, each a key", name, len(f.Choices))200 }201 return f, nil202}203204// Parse reads what the person typed for a Text, Number or Integer205// field.206func (f Field) Parse(text string) (any, error) {207 switch f.Kind {208 case Number:209 n, err := strconv.ParseFloat(strings.TrimSpace(text), 64)210 if err != nil {211 return nil, fmt.Errorf("%s takes a number", f.Label())212 }213 return n, nil214 case Integer:215 n, err := strconv.Atoi(strings.TrimSpace(text))216 if err != nil {217 return nil, fmt.Errorf("%s takes a whole number", f.Label())218 }219 return n, nil220 }221 return text, nil222}223224// ChoiceLabel is how choice i is shown.225func (f Field) ChoiceLabel(i int) string {226 if i < len(f.Labels) && f.Labels[i] != "" {227 return f.Labels[i]228 }229 return f.Choices[i]230}231232// AddAsk puts a question on the board for a person to answer.233func (b *Board) AddAsk(a *Ask) {234 b.nextAsk++235 a.ID = b.nextAsk236 b.asks = append(b.asks, a)237}238239// NextAsk is the first question waiting that the person at seat should240// answer: the host's, or in a room without one, anyone's not already241// shown to another who is here (present). It is then seat's to show.242func (b *Board) NextAsk(seat int, present func(id int) bool) *Ask {243 if b.Host != 0 && b.Host != seat && present(b.Host) {244 return nil245 }246 for _, a := range b.asks {247 if a.taken == seat || a.taken == 0 || !present(a.taken) {248 a.taken = seat249 return a250 }251 }252 return nil253}254255// Answer answers a question and takes it off the board.256func (b *Board) Answer(a *Ask, ans Answer) {257 if a.done {258 return259 }260 a.done = true261 b.asks = slices.DeleteFunc(b.asks, func(x *Ask) bool { return x == a })262 if a.Kind == Grant && ans.Action == "accept" && ans.Always {263 switch a.Grant {264 case "notebooks":265 b.Grants.Notebooks = true266 case "jev":267 b.Grants.JEV = true268 }269 }270 a.reply <- ans271}272273// Withdraw takes a question off the board unanswered: the agent gave274// up waiting, or left.275func (b *Board) Withdraw(a *Ask) {276 if a.done {277 return278 }279 a.done = true280 b.asks = slices.DeleteFunc(b.asks, func(x *Ask) bool { return x == a })281}282283// Asks are the questions waiting, oldest first.284func (b *Board) Asks() []*Ask { return slices.Clone(b.asks) }