1package subagents
2
3import (
4 "bufio"
5 "bytes"
6 "encoding/json"
7 "errors"
8 "fmt"
9 "io"
10 "strings"
11 "time"
12
13 "github.com/google/uuid"
14)
15
16// ProtocolVersion is the version of the supervisor/worker JSONL protocol.
17const ProtocolVersion = 1
18
19// Canonical parent-to-worker command names.
20const (
21 CommandAgentPing = "agent.ping"
22 CommandTurnStart = "turn.start"
23 CommandTurnCancel = "turn.cancel"
24 CommandAgentShutdown = "agent.shutdown"
25)
26
27// Canonical worker-to-parent event names.
28const (
29 EventAgentReady = "agent.ready"
30 EventAgentHeartbeat = "agent.heartbeat"
31 EventTurnStarted = "turn.started"
32 EventTurnProgress = "turn.progress"
33 EventToolStarted = "tool.started"
34 EventToolFinished = "tool.finished"
35 EventMessageDelta = "message.delta"
36 EventTurnResult = "turn.result"
37 EventTurnFailed = "turn.failed"
38 EventAgentIdle = "agent.idle"
39 EventAgentExited = "agent.exited"
40)
41
42var (
43 // ErrEmptyProtocolLine is returned for a blank JSONL line or command.
44 ErrEmptyProtocolLine = errors.New("subagent protocol: empty line")
45 // ErrNotCommand is returned when ParseCommand is given a JSON event.
46 ErrNotCommand = errors.New("subagent protocol: message is not a command")
47 // ErrNotEvent is returned when ParseEvent is given a JSON command.
48 ErrNotEvent = errors.New("subagent protocol: message is not an event")
49 // ErrUnsupportedProtocolVersion is returned by Validate for a version
50 // this implementation cannot interpret.
51 ErrUnsupportedProtocolVersion = errors.New("subagent protocol: unsupported version")
52)
53
54var commandNames = map[string]struct{}{
55 CommandAgentPing: {},
56 CommandTurnStart: {},
57 CommandTurnCancel: {},
58 CommandAgentShutdown: {},
59}
60
61var eventNames = map[string]struct{}{
62 EventAgentReady: {},
63 EventAgentHeartbeat: {},
64 EventTurnStarted: {},
65 EventTurnProgress: {},
66 EventToolStarted: {},
67 EventToolFinished: {},
68 EventMessageDelta: {},
69 EventTurnResult: {},
70 EventTurnFailed: {},
71 EventAgentIdle: {},
72 EventAgentExited: {},
73}
74
75// IsCommandName reports whether name is one of the commands defined by this
76// version of the protocol. Readers intentionally do not use this as a parse
77// gate: a newer command can still be decoded and forwarded by an older
78// supervisor.
79func IsCommandName(name string) bool {
80 _, ok := commandNames[name]
81 return ok
82}
83
84// IsEventName reports whether name is one of the events defined by this
85// version of the protocol. Unknown event names are valid protocol messages.
86func IsEventName(name string) bool {
87 _, ok := eventNames[name]
88 return ok
89}
90
91// MessageID returns a fresh opaque id suitable for correlating a command and
92// its result. UUIDs are used rather than a counter so ids remain unique when
93// several supervisors or workers are running at once.
94func NewMessageID() string { return uuid.NewString() }
95
96// IsMessageID reports whether id has the UUID form emitted by NewMessageID.
97// Message ids are otherwise opaque on the wire, so readers do not require
98// this form when accepting messages from a newer implementation.
99func IsMessageID(id string) bool {
100 if id == "" {
101 return false
102 }
103 _, err := uuid.Parse(id)
104 return err == nil
105}
106
107// Envelope is one versioned JSONL command or event. Payload is deliberately a
108// RawMessage: an old supervisor can retain an event it does not know about,
109// including fields added inside that event's payload, and write it back
110// without decoding it into a lossy map.
111//
112// The JSON representation of a current message is:
113//
114// {"version":1,"type":"...","message_id":"...","agent_id":"...",
115// "turn_id":"...","timestamp":"...","payload":{...}}
116//
117// TurnID is omitted when it is empty. Unknown top-level fields are retained in
118// Unknown.
119type Envelope struct {
120 Version int `json:"version"`
121 Type string `json:"type"`
122 MessageID string `json:"message_id,omitempty"`
123 AgentID string `json:"agent_id,omitempty"`
124 TurnID string `json:"turn_id,omitempty"`
125 Timestamp time.Time `json:"timestamp"`
126 Payload json.RawMessage `json:"payload"`
127
128 // Unknown contains unrecognised top-level fields from a versioned
129 // envelope. Payload itself retains unrecognised payload fields.
130 Unknown map[string]json.RawMessage `json:"-"`
131
132 // constructionErr records a payload-construction failure for the
133 // compatibility-preserving value constructor. It is surfaced by Validate
134 // and MarshalJSON rather than being turned into a valid null payload.
135 constructionErr error
136}
137
138// NewEnvelope creates a current-protocol envelope and assigns a message id
139// and timestamp. payload may be any JSON-marshalable value. A nil payload is
140// encoded as an empty object, which keeps command and event payloads uniform.
141//
142// For a marshalable payload this helper cannot fail. Call
143// NewEnvelopeWithPayload when construction errors must be reported instead of
144// being represented as a null payload.
145func NewEnvelope(typ, agentID, turnID string, payload ...any) Envelope {
146 e, err := NewEnvelopeWithPayload(typ, agentID, turnID, firstPayload(payload))
147 if err == nil {
148 return e
149 }
150 // Preserve the value-returning API used by command/event call sites, but
151 // retain the construction failure so validation and wire encoding fail
152 // instead of silently sending a null payload.
153 return Envelope{
154 Version: ProtocolVersion,
155 Type: typ,
156 MessageID: NewMessageID(),
157 AgentID: agentID,
158 TurnID: turnID,
159 Timestamp: time.Now().UTC(),
160 Payload: nil,
161 constructionErr: fmt.Errorf("subagent protocol payload: %w", err),
162 }
163}
164
165// NewEnvelopeWithPayload is the checked form of NewEnvelope.
166func NewEnvelopeWithPayload(typ, agentID, turnID string, payload any) (Envelope, error) {
167 if strings.TrimSpace(typ) == "" {
168 return Envelope{}, errors.New("subagent protocol: empty message type")
169 }
170 p, err := marshalPayload(payload)
171 if err != nil {
172 return Envelope{}, fmt.Errorf("subagent protocol payload: %w", err)
173 }
174 return Envelope{
175 Version: ProtocolVersion,
176 Type: typ,
177 MessageID: NewMessageID(),
178 AgentID: agentID,
179 TurnID: turnID,
180 Timestamp: time.Now().UTC(),
181 Payload: p,
182 }, nil
183}
184
185// NewCommand creates a current-protocol command. Unknown command names are
186// allowed so a bridge can forward commands introduced by a newer peer.
187func NewCommand(name, agentID, turnID string, payload ...any) Envelope {
188 return NewEnvelope(name, agentID, turnID, payload...)
189}
190
191// NewEventEnvelope creates a current-protocol event. The name is not checked
192// against the known event registry so unknown future events remain forward-compatible.
193func NewEventEnvelope(name, agentID, turnID string, payload ...any) Envelope {
194 return NewEnvelope(name, agentID, turnID, payload...)
195}
196
197func firstPayload(payload []any) any {
198 if len(payload) == 0 || payload[0] == nil {
199 return map[string]any{}
200 }
201 return payload[0]
202}
203
204func marshalPayload(payload any) (json.RawMessage, error) {
205 if payload == nil {
206 return json.RawMessage("{}"), nil
207 }
208 b, err := json.Marshal(payload)
209 if err != nil {
210 return nil, err
211 }
212 return append(json.RawMessage(nil), b...), nil
213}
214
215// Validate checks the required fields of a current-protocol envelope. It is
216// intentionally separate from ParseEnvelope: parsing an unknown future event
217// should succeed, while a sender can opt into strict validation before write.
218func (e Envelope) Validate() error {
219 if e.constructionErr != nil {
220 return e.constructionErr
221 }
222 if e.Version != ProtocolVersion {
223 return fmt.Errorf("%w: %d", ErrUnsupportedProtocolVersion, e.Version)
224 }
225 if e.Type == "" {
226 return errors.New("subagent protocol: missing type")
227 }
228 if e.MessageID == "" {
229 return errors.New("subagent protocol: missing message_id")
230 }
231 if e.AgentID == "" {
232 return errors.New("subagent protocol: missing agent_id")
233 }
234 if e.Timestamp.IsZero() {
235 return errors.New("subagent protocol: missing timestamp")
236 }
237 if len(e.Payload) == 0 || !json.Valid(e.Payload) {
238 return errors.New("subagent protocol: invalid payload")
239 }
240 return nil
241}
242
243// IsCommand reports whether this envelope has a known command name.
244func (e Envelope) IsCommand() bool { return IsCommandName(e.Type) }
245
246// IsEvent reports whether this envelope has a known event name.
247func (e Envelope) IsEvent() bool { return IsEventName(e.Type) }
248
249// DecodePayload decodes the payload into dst without changing the preserved
250// raw payload. This is the preferred way for a consumer to inspect a known
251// event while still retaining it for a later forward-compatible write.
252func (e Envelope) DecodePayload(dst any) error {
253 if dst == nil {
254 return errors.New("subagent protocol: nil payload destination")
255 }
256 if len(e.Payload) == 0 {
257 return errors.New("subagent protocol: missing payload")
258 }
259 return json.Unmarshal(e.Payload, dst)
260}
261
262// PayloadFields returns an object payload as raw fields. It returns an error
263// for scalar or array payloads, rather than silently discarding their shape.
264func (e Envelope) PayloadFields() (map[string]json.RawMessage, error) {
265 var fields map[string]json.RawMessage
266 if err := e.DecodePayload(&fields); err != nil {
267 return nil, err
268 }
269 if fields == nil {
270 return nil, errors.New("subagent protocol: payload is not an object")
271 }
272 return fields, nil
273}
274
275// PayloadValue returns a decoded generic payload. PayloadFields is preferable
276// when the caller needs to preserve exact nested JSON values.
277func (e Envelope) PayloadValue() (any, error) {
278 var value any
279 if err := e.DecodePayload(&value); err != nil {
280 return nil, err
281 }
282 return value, nil
283}
284
285// SetPayload replaces the payload while keeping the envelope metadata and
286// message id intact.
287func (e *Envelope) SetPayload(payload any) error {
288 if e == nil {
289 return errors.New("subagent protocol: nil envelope")
290 }
291 p, err := marshalPayload(payload)
292 if err != nil {
293 return fmt.Errorf("subagent protocol payload: %w", err)
294 }
295 e.Payload = p
296 return nil
297}
298
299// ParseEnvelope parses one current-protocol envelope. Unknown event types and
300// unknown payload fields are valid and retained.
301func ParseEnvelope(data []byte) (Envelope, error) {
302 data = bytes.TrimSpace(data)
303 if len(data) == 0 {
304 return Envelope{}, ErrEmptyProtocolLine
305 }
306 var raw map[string]json.RawMessage
307 if err := json.Unmarshal(data, &raw); err != nil {
308 return Envelope{}, fmt.Errorf("subagent protocol: decode envelope: %w", err)
309 }
310 if raw == nil {
311 return Envelope{}, errors.New("subagent protocol: envelope must be an object")
312 }
313 var e Envelope
314 if err := unmarshalVersionedEnvelope(raw, &e); err != nil {
315 return Envelope{}, err
316 }
317 return e, nil
318}
319
320// ParseJSONL parses one newline-delimited JSON message. It accepts a final
321// line without a newline, as is customary for a stream after a clean EOF.
322func ParseJSONL(line []byte) (Envelope, error) { return ParseEnvelope(line) }
323
324// ParseEvent parses a current-protocol JSON event. A known command is
325// rejected; unknown types are accepted because an older peer cannot know
326// whether a future event name is present in its event registry.
327func ParseEvent(data []byte) (Envelope, error) {
328 e, err := ParseEnvelope(data)
329 if err != nil {
330 return Envelope{}, err
331 }
332 if e.IsCommand() {
333 return Envelope{}, ErrNotEvent
334 }
335 return e, nil
336}
337
338// ParseCommand parses one current-protocol JSON command.
339func ParseCommand(line string) (Envelope, error) {
340 raw := strings.TrimSpace(line)
341 if raw == "" {
342 return Envelope{}, ErrEmptyProtocolLine
343 }
344 e, err := ParseEnvelope([]byte(raw))
345 if err != nil {
346 return Envelope{}, err
347 }
348 if IsEventName(e.Type) {
349 return Envelope{}, ErrNotCommand
350 }
351 return e, nil
352}
353
354// MarshalEnvelope returns one JSON object without a trailing newline.
355func MarshalEnvelope(e Envelope) ([]byte, error) {
356 return json.Marshal(e)
357}
358
359// MarshalJSONL returns one newline-terminated JSONL message.
360func MarshalJSONL(e Envelope) ([]byte, error) {
361 b, err := MarshalEnvelope(e)
362 if err != nil {
363 return nil, err
364 }
365 return append(b, '\n'), nil
366}
367
368// WriteEnvelope writes one complete newline-terminated JSONL message.
369func WriteEnvelope(w io.Writer, e Envelope) error {
370 if w == nil {
371 return errors.New("subagent protocol: nil writer")
372 }
373 b, err := MarshalJSONL(e)
374 if err != nil {
375 return err
376 }
377 n, err := w.Write(b)
378 if err == nil && n != len(b) {
379 return io.ErrShortWrite
380 }
381 return err
382}
383
384// ReadEnvelope reads exactly one JSONL message from a buffered reader. A
385// caller reading many messages should reuse the same bufio.Reader.
386func ReadEnvelope(r *bufio.Reader) (Envelope, error) {
387 if r == nil {
388 return Envelope{}, errors.New("subagent protocol: nil reader")
389 }
390 line, err := r.ReadBytes('\n')
391 if len(bytes.TrimSpace(line)) == 0 && err == io.EOF {
392 return Envelope{}, io.EOF
393 }
394 if len(line) != 0 {
395 parsed, parseErr := ParseJSONL(line)
396 if parseErr != nil {
397 return Envelope{}, parseErr
398 }
399 return parsed, nil
400 }
401 return Envelope{}, err
402}
403
404func (e Envelope) MarshalJSON() ([]byte, error) {
405 return e.marshalVersionedJSON()
406}
407
408func (e Envelope) marshalVersionedJSON() ([]byte, error) {
409 if e.constructionErr != nil {
410 return nil, e.constructionErr
411 }
412 version := e.Version
413 if version == 0 {
414 version = ProtocolVersion
415 }
416 messageID := e.MessageID
417 if messageID == "" {
418 messageID = NewMessageID()
419 }
420 timestamp := e.Timestamp
421 if timestamp.IsZero() {
422 timestamp = time.Now().UTC()
423 }
424 payload := e.Payload
425 if len(payload) == 0 {
426 payload = json.RawMessage("{}")
427 }
428 if !json.Valid(payload) {
429 return nil, errors.New("subagent protocol: invalid payload JSON")
430 }
431 out := make(map[string]json.RawMessage, len(e.Unknown)+7)
432 for key, value := range e.Unknown {
433 if isEnvelopeField(key) {
434 continue
435 }
436 out[key] = cloneRaw(value)
437 }
438 putRaw(out, "version", version)
439 putRaw(out, "type", e.Type)
440 putRaw(out, "message_id", messageID)
441 putRaw(out, "agent_id", e.AgentID)
442 if e.TurnID != "" {
443 putRaw(out, "turn_id", e.TurnID)
444 }
445 putRaw(out, "timestamp", timestamp)
446 out["payload"] = cloneRaw(payload)
447 return json.Marshal(out)
448}
449
450func (e *Envelope) UnmarshalJSON(data []byte) error {
451 var raw map[string]json.RawMessage
452 if err := json.Unmarshal(data, &raw); err != nil {
453 return err
454 }
455 if raw == nil {
456 return errors.New("subagent protocol: envelope must be an object")
457 }
458 var parsed Envelope
459 if err := unmarshalVersionedEnvelope(raw, &parsed); err != nil {
460 return err
461 }
462 *e = parsed
463 return nil
464}
465
466func unmarshalVersionedEnvelope(raw map[string]json.RawMessage, e *Envelope) error {
467 if err := decodeRequired(raw, "version", &e.Version); err != nil {
468 return err
469 }
470 if e.Version != ProtocolVersion {
471 return fmt.Errorf("%w: %d", ErrUnsupportedProtocolVersion, e.Version)
472 }
473 if err := decodeRequired(raw, "type", &e.Type); err != nil {
474 return err
475 }
476 if err := decodeOptionalString(raw, "message_id", &e.MessageID); err != nil {
477 return err
478 }
479 if err := decodeOptionalString(raw, "agent_id", &e.AgentID); err != nil {
480 return err
481 }
482 if err := decodeOptionalString(raw, "turn_id", &e.TurnID); err != nil {
483 return err
484 }
485 if timestamp, ok := raw["timestamp"]; ok {
486 if err := json.Unmarshal(timestamp, &e.Timestamp); err != nil {
487 return fmt.Errorf("subagent protocol: timestamp: %w", err)
488 }
489 }
490 if payload, ok := raw["payload"]; ok {
491 e.Payload = cloneRaw(payload)
492 } else {
493 e.Payload = json.RawMessage("{}")
494 }
495 e.Unknown = unknownFields(raw, map[string]struct{}{
496 "version": {}, "type": {}, "message_id": {}, "agent_id": {},
497 "turn_id": {}, "timestamp": {}, "payload": {},
498 })
499 return nil
500}
501
502func decodeRequired(raw map[string]json.RawMessage, name string, dst any) error {
503 value, ok := raw[name]
504 if !ok {
505 return fmt.Errorf("subagent protocol: missing %s", name)
506 }
507 if err := json.Unmarshal(value, dst); err != nil {
508 return fmt.Errorf("subagent protocol: %s: %w", name, err)
509 }
510 if name == "type" {
511 if text, ok := dst.(*string); ok && strings.TrimSpace(*text) == "" {
512 return errors.New("subagent protocol: empty message type")
513 }
514 }
515 return nil
516}
517
518func decodeOptionalString(raw map[string]json.RawMessage, name string, dst *string) error {
519 value, ok := raw[name]
520 if !ok {
521 return nil
522 }
523 if err := json.Unmarshal(value, dst); err != nil {
524 return fmt.Errorf("subagent protocol: %s: %w", name, err)
525 }
526 return nil
527}
528
529func unknownFields(raw map[string]json.RawMessage, known map[string]struct{}) map[string]json.RawMessage {
530 var out map[string]json.RawMessage
531 for key, value := range raw {
532 if _, ok := known[key]; ok {
533 continue
534 }
535 if out == nil {
536 out = make(map[string]json.RawMessage)
537 }
538 out[key] = cloneRaw(value)
539 }
540 return out
541}
542
543func isEnvelopeField(key string) bool {
544 switch key {
545 case "version", "type", "message_id", "agent_id", "turn_id", "timestamp", "payload":
546 return true
547 default:
548 return false
549 }
550}
551
552func putRaw(out map[string]json.RawMessage, key string, value any) {
553 b, err := json.Marshal(value)
554 if err == nil {
555 out[key] = b
556 }
557}
558
559func cloneRaw(value json.RawMessage) json.RawMessage {
560 return append(json.RawMessage(nil), value...)
561}
562
563// The following small payload types document the stable fields of the
564// protocol without forcing a receiver to use them. Receivers can always use
565// Envelope.Payload to retain fields added in later versions.
566type AgentPingPayload struct{}
567
568type ShutdownOrigin string
569
570const (
571 ShutdownOriginTargeted ShutdownOrigin = "targeted"
572 ShutdownOriginSession ShutdownOrigin = "session"
573 ShutdownOriginDeadline ShutdownOrigin = "deadline"
574 ShutdownOriginProcess ShutdownOrigin = "process"
575)
576
577func (o ShutdownOrigin) Sanitized() ShutdownOrigin {
578 switch o {
579 case ShutdownOriginTargeted, ShutdownOriginSession, ShutdownOriginDeadline, ShutdownOriginProcess:
580 return o
581 default:
582 return ""
583 }
584}
585
586type AgentShutdownPayload struct {
587 Origin ShutdownOrigin `json:"origin,omitempty"`
588}
589type AgentReadyPayload struct {
590 // Version is the worker/application version. WorkerVersion is the
591 // unambiguous spelling for new producers.
592 Version string `json:"version,omitempty"`
593 WorkerVersion string `json:"worker_version,omitempty"`
594 CWD string `json:"cwd,omitempty"`
595 Model string `json:"model,omitempty"`
596 Provider string `json:"provider,omitempty"`
597}
598type AgentHeartbeatPayload struct {
599 Activity string `json:"activity,omitempty"`
600}
601type TurnStartPayload struct {
602 Prompt string `json:"prompt"`
603 // NewRun starts a fresh max_turns budget while retaining the worker's
604 // existing session and transcript. It is set only for explicit
605 // subagent_resume follow-ups.
606 NewRun bool `json:"new_run,omitempty"`
607}
608type TurnCancelPayload struct {
609 Reason string `json:"reason,omitempty"`
610}
611type TurnProgressPayload struct {
612 Text string `json:"text,omitempty"`
613 Percent float64 `json:"percent,omitempty"`
614 CurrentTool string `json:"current_tool,omitempty"`
615}
616type ToolStartedPayload struct {
617 ToolID string `json:"tool_id,omitempty"`
618 Name string `json:"name"`
619}
620type ToolFinishedPayload struct {
621 ToolID string `json:"tool_id,omitempty"`
622 Name string `json:"name,omitempty"`
623 Error string `json:"error,omitempty"`
624}
625type MessageDeltaPayload struct {
626 Text string `json:"text,omitempty"`
627}
628type ArtifactReference struct {
629 Name string `json:"name,omitempty"`
630 Ref string `json:"ref"`
631 Mime string `json:"mime,omitempty"`
632 Size int64 `json:"size,omitempty"`
633}
634type ProtocolError struct {
635 Code string `json:"code,omitempty"`
636 Message string `json:"message,omitempty"`
637}
638type TurnResultPayload struct {
639 Status string `json:"status"`
640 Summary string `json:"summary,omitempty"`
641 Output string `json:"output,omitempty"`
642 Structured json.RawMessage `json:"structured,omitempty"`
643 Artifacts []ArtifactReference `json:"artifacts,omitempty"`
644 ChangedFiles []string `json:"changed_files,omitempty"`
645 Usage map[string]any `json:"usage,omitempty"`
646 Error *ProtocolError `json:"error"`
647}
648type TurnFailedPayload struct {
649 Error *ProtocolError `json:"error,omitempty"`
650}
651type AgentExitedPayload struct {
652 Code int `json:"code,omitempty"`
653 Reason string `json:"reason,omitempty"`
654 Origin ShutdownOrigin `json:"origin,omitempty"`
655}