byre/internal/agents/inbox.go

package agents

import (
	"context"
	"encoding/json"
	"errors"
	"fmt"
	"log/slog"
	"slices"
	"strings"
	"time"

	"g.fc.run/misc/byre/internal/api"
	"g.fc.run/misc/byre/internal/mcp"
	"g.fc.run/misc/byre/internal/store"
)

var attachmentKinds = []api.AttachmentKind{
	api.AttachmentKindFile, api.AttachmentKindDiff, api.AttachmentKindImage, api.AttachmentKindDiagram,
	api.AttachmentKindTask, api.AttachmentKindSnapshot, api.AttachmentKindUrl,
}

type askArgs struct {
	Context     string           `json:"context"`
	Question    string           `json:"question"`
	WhyBlocked  string           `json:"why_blocked"`
	Attachments []api.Attachment `json:"attachments"`
	Blocking    *bool            `json:"blocking"`
	ThreadID    string           `json:"thread_id"`
}

func (m *Manager) askTool(l *live) mcp.Tool {
	return mcp.Tool{
		Name: "ask_human",
		Description: "Ask the human a question through their inbox. Use this for every question; never ask in plain prose. " +
			"The human answers from their inbox without having read your transcript, so every question must stand on its own: " +
			"put the background they need in context. " +
			"Write the choices inline in context as Markdown list items, each with what it would concretely do: " +
			"`- ( ) choice` for pick-one choices (mark the one you recommend `- (*) choice`), or " +
			"`- [ ] choice` for pick-any choices (mark the ones you recommend `- [x] choice`). " +
			"The human picks them where they appear; there is no separate list of options. " +
			"If blocking is true, end your turn after asking; " +
			"the answer arrives as your next message. If blocking is false, continue with work that doesn't depend on the answer. " +
			"Ask several questions in one turn by calling this once per question; they are answered together. " +
			"To answer a follow-up question from the human, call this again with the thread_id they gave you.",
		InputSchema: json.RawMessage(`{
			"type": "object",
			"properties": {
				"context": {"type": "string", "description": "Everything the human needs to answer well without reading your transcript: what you are working on, what you found, the relevant facts and data, and the choices. Several paragraphs is normal. Markdown is rendered. Each choice is a list item on its own line, at least two, all of one kind: '- ( ) Merge now: ships the fix today' / '- (*) Wait for review: ...' (pick one; exactly one (*) marks your recommendation), or '- [ ] Fix the tests' / '- [x] Update the docs' (pick any; [x] marks the ones you recommend). The choice's text is the whole line after the marker, and it comes back verbatim as the answer. Put long supporting detail (logs, full data, per-option analysis) in collapsed sections the human expands on demand: a line '<details><summary>Title</summary>', the Markdown content, then a line '</details>'."},
				"question": {"type": "string", "description": "The question itself, in one or two sentences. Markdown is rendered."},
				"why_blocked": {"type": "string", "description": "What you can't do until this is answered, and why you can't decide it yourself."},
				"attachments": {"type": "array", "items": {
					"type": "object",
					"properties": {
						"kind": {"type": "string", "enum": ["file", "diff", "image", "diagram", "task", "snapshot", "url"]},
						"path": {"type": "string"},
						"start_line": {"type": "integer"},
						"end_line": {"type": "integer"},
						"ref": {"type": "string"},
						"url": {"type": "string"},
						"mermaid": {"type": "string"},
						"task_id": {"type": "string"},
						"snapshot_id": {"type": "string"},
						"caption": {"type": "string"}
					},
					"required": ["kind"]
				}},
				"blocking": {"type": "boolean"},
				"thread_id": {"type": "string"}
			},
			"required": ["context", "question", "why_blocked", "blocking"]
		}`),
		Handler: func(ctx context.Context, raw json.RawMessage) (string, error) {
			return m.ask(ctx, l, raw)
		},
	}
}

func (m *Manager) ask(ctx context.Context, l *live, raw json.RawMessage) (string, error) {
	args := askArgs{}
	err := json.Unmarshal(raw, &args)
	if err != nil {
		return "", fmt.Errorf("malformed ask_human call: %w", err)
	}

	err = m.validateAsk(ctx, l, args)
	if err != nil {
		return "", err
	}

	id := store.NewID()
	thread := args.ThreadID
	if thread == "" {
		thread = id
	}

	priority := api.Fyi
	if *args.Blocking {
		priority = api.Blocked
	}

	attachments := args.Attachments
	if attachments == nil {
		attachments = []api.Attachment{}
	}

	choices := parseChoices(args.Context)
	options := []string{}
	recommended := []string{}
	for _, c := range choices {
		options = append(options, c.text)
		if c.recommended {
			recommended = append(recommended, c.text)
		}
	}

	payload := map[string]any{
		"context":     args.Context,
		"question":    args.Question,
		"why_blocked": args.WhyBlocked,
		"options":     options,
		"multiple":    choices[0].checkbox,
		"recommended": recommended,
		"attachments": attachments,
	}

	err = m.store.CreateInboxItem(ctx, m.newItem(l, id, thread, api.Question, priority, *args.Blocking, store.ToPayload(payload)))
	if err != nil {
		return "", err
	}

	m.mu.Lock()
	l.asked = true
	l.askedBlock = l.askedBlock || *args.Blocking
	m.mu.Unlock()

	cancel := fmt.Sprintf("If the human answers it some other way, such as in a chat message, withdraw it with cancel_question and question_id %q.", id)
	if *args.Blocking {
		return fmt.Sprintf("Question filed in the human's inbox. question_id: %s. It is blocking: end your turn now. The answer will arrive as your next message. %s", id, cancel), nil
	}
	return fmt.Sprintf("Question filed in the human's inbox. question_id: %s. Continue with work that doesn't depend on it; the answer will arrive as a later message. %s", id, cancel), nil
}

type cancelArgs struct {
	QuestionID string `json:"question_id"`
}

func (m *Manager) cancelQuestionTool(l *live) mcp.Tool {
	return mcp.Tool{
		Name: "cancel_question",
		Description: "Withdraw a question you asked with ask_human that is still waiting in the human's inbox, " +
			"for example because the human answered it in a chat message instead. It is removed from the inbox and no answer to it will arrive.",
		InputSchema: json.RawMessage(`{
			"type": "object",
			"properties": {
				"question_id": {"type": "string", "description": "The question_id ask_human returned."}
			},
			"required": ["question_id"]
		}`),
		Handler: func(ctx context.Context, raw json.RawMessage) (string, error) {
			return m.cancelQuestion(ctx, l, raw)
		},
	}
}

func (m *Manager) cancelQuestion(ctx context.Context, l *live, raw json.RawMessage) (string, error) {
	args := cancelArgs{}
	err := json.Unmarshal(raw, &args)
	if err != nil {
		return "", fmt.Errorf("malformed cancel_question call: %w", err)
	}

	item, err := m.store.GetInboxItem(ctx, args.QuestionID)
	if errors.Is(err, store.ErrNotFound) {
		return "", fmt.Errorf("question_id %s does not exist", args.QuestionID)
	}
	if err != nil {
		return "", err
	}
	if item.AgentSessionId != l.id || item.Kind != api.Question {
		return "", fmt.Errorf("question_id %s is not a question you asked", args.QuestionID)
	}

	err = m.store.WithdrawItem(ctx, item.Id)
	if errors.Is(err, store.ErrNotOpen) {
		return "", fmt.Errorf("question %s is already %s; nothing was cancelled", item.Id, item.State)
	}
	if err != nil {
		return "", err
	}
	return fmt.Sprintf("Question %s withdrawn from the human's inbox.", item.Id), nil
}

func (m *Manager) validateAsk(ctx context.Context, l *live, args askArgs) error {
	problems := []string{}

	if strings.TrimSpace(args.Context) == "" {
		problems = append(problems, "context is empty: give the background the human needs to answer without reading your transcript")
	}
	if strings.TrimSpace(args.Question) == "" {
		problems = append(problems, "question is empty")
	}
	if strings.TrimSpace(args.WhyBlocked) == "" {
		problems = append(problems, "why_blocked is empty: say what you can't do until this is answered and why you can't decide it yourself")
	}
	choices := parseChoices(args.Context)
	if len(choices) < 2 {
		problems = append(problems, "context needs at least two choices, each a list item on its own line: `- ( ) choice` (pick one) or `- [ ] choice` (pick any)")
	}
	seen := map[string]bool{}
	radios := 0
	recommendedRadios := 0
	for _, c := range choices {
		if seen[c.text] {
			problems = append(problems, fmt.Sprintf("choice %q appears twice; each choice's text must be unique", c.text))
		}
		seen[c.text] = true
		if !c.checkbox {
			radios++
			if c.recommended {
				recommendedRadios++
			}
		}
	}
	if radios > 0 && radios < len(choices) {
		problems = append(problems, "choices mix `( )` and `[ ]`: use one kind per question")
	}
	if radios == len(choices) && len(choices) > 0 && recommendedRadios != 1 {
		problems = append(problems, "mark exactly one pick-one choice as recommended with `- (*) choice`")
	}
	if args.Blocking == nil {
		problems = append(problems, "blocking is required")
	}
	for i, a := range args.Attachments {
		if !slices.Contains(attachmentKinds, a.Kind) {
			problems = append(problems, fmt.Sprintf("attachments[%d].kind %q is not one of %v", i, a.Kind, attachmentKinds))
		}
	}

	if args.ThreadID != "" {
		item, err := m.store.GetInboxItem(ctx, args.ThreadID)
		switch {
		case errors.Is(err, store.ErrNotFound):
			problems = append(problems, fmt.Sprintf("thread_id %s does not exist", args.ThreadID))
		case err != nil:
			return err
		case item.AgentSessionId != l.id:
			problems = append(problems, fmt.Sprintf("thread_id %s belongs to another session", args.ThreadID))
		}
	}

	if len(problems) > 0 {
		return fmt.Errorf("ask_human rejected, nothing was filed. Fix and call again: %s", strings.Join(problems, "; "))
	}
	return nil
}

func (m *Manager) fileSummary(l *live, text string) string {
	id := store.NewID()
	err := m.store.CreateInboxItem(context.Background(), m.newItem(l, id, id, api.Report, api.Fyi, false, map[string]any{"text": text}))
	if err != nil {
		slog.Error("file turn summary", "session", l.id, "err", err)
		return ""
	}
	return id
}

func (m *Manager) fileError(l *live, msg string) {
	id := store.NewID()
	err := m.store.CreateInboxItem(context.Background(), m.newItem(l, id, id, api.SessionError, api.Blocked, true, map[string]any{"error": msg}))
	if err != nil {
		slog.Error("file session error", "session", l.id, "err", err)
	}
}

func (m *Manager) fileRateLimited(l *live, info map[string]any, failure string) {
	id := store.NewID()
	payload := map[string]any{"vendor": string(l.vendor), "rate_limit": info, "error": failure}
	err := m.store.CreateInboxItem(context.Background(), m.newItem(l, id, id, api.RateLimited, api.Blocked, true, payload))
	if err != nil {
		slog.Error("file rate limit", "session", l.id, "err", err)
	}
}

func (m *Manager) newItem(l *live, id, thread string, kind api.InboxItemKind, priority api.InboxItemPriority, blocking bool, payload map[string]any) api.InboxItem {
	return api.InboxItem{
		Id:             id,
		ProjectId:      l.projectID,
		TaskId:         l.taskID,
		AgentSessionId: l.id,
		ThreadId:       thread,
		Kind:           kind,
		Priority:       priority,
		Blocking:       blocking,
		Payload:        payload,
		State:          api.Open,
		CreatedAt:      time.Now().UTC(),
		Responses:      []api.InboxResponse{},
	}
}

func (m *Manager) Deliver(ctx context.Context, sessionID string) error {
	m.deliverMu.Lock()
	defer m.deliverMu.Unlock()

	pending, err := m.store.Undelivered(ctx, sessionID)
	if err != nil {
		return err
	}
	if len(pending) == 0 {
		return nil
	}
	ids := responseIDs(pending)

	sess, err := m.store.GetSession(ctx, sessionID)
	if err != nil {
		return err
	}
	if sess.State != api.AgentSessionStateIdle {
		slog.Info("deliver held: session not idle", "session", sessionID, "state", sess.State, "responses", ids)
		return nil
	}

	waiting, err := m.store.AwaitingHuman(ctx, sessionID)
	if err != nil {
		return err
	}
	if waiting > 0 {
		slog.Info("deliver held: blocking question waiting", "session", sessionID, "waiting", waiting, "responses", ids)
		return nil
	}

	task, err := m.store.GetTask(ctx, sess.TaskId)
	if err != nil {
		return err
	}

	after, err := m.turn(ctx, sess, task, pending[0].Response.User, deliveryMessage(pending), nil)
	if errors.Is(err, ErrConflict) {
		slog.Info("deliver held: session changed state", "session", sessionID, "err", err, "responses", ids)
		return nil
	}
	if err != nil {
		slog.Error("deliver failed", "session", sessionID, "err", err, "responses", ids)
		return err
	}

	m.mu.Lock()
	l := m.live[sessionID]
	closing := m.closing
	m.mu.Unlock()
	pid := 0
	exited := false
	if l != nil {
		pid = l.sess.Pid()
		exited = l.sess.Exited()
	}
	slog.Info("deliver sent", "session", sessionID, "state", after.State, "pid", pid, "exited", exited, "closing", closing, "responses", ids)
	return m.store.MarkDelivered(ctx, ids)
}

func (m *Manager) deliverMidTurn(ctx context.Context, l *live) string {
	m.deliverMu.Lock()
	defer m.deliverMu.Unlock()

	pending, err := m.store.Undelivered(ctx, l.id)
	if err != nil {
		slog.Error("list undelivered responses", "session", l.id, "err", err)
		return ""
	}
	if len(pending) == 0 {
		return ""
	}
	ids := responseIDs(pending)

	waiting, err := m.store.AwaitingHuman(ctx, l.id)
	if err != nil {
		slog.Error("count blocking questions", "session", l.id, "err", err)
		return ""
	}
	if waiting > 0 {
		slog.Info("deliver held mid-turn: blocking question waiting", "session", l.id, "waiting", waiting, "responses", ids)
		return ""
	}

	err = m.store.MarkDelivered(ctx, ids)
	if err != nil {
		slog.Error("mark responses delivered", "session", l.id, "err", err, "responses", ids)
		return ""
	}

	text := deliveryMessage(pending)
	err = m.emit(ctx, l, pending[0].Response.User, "user.message", map[string]any{"text": text, "mid_turn": true})
	if err != nil {
		slog.Error("record mid-turn delivery", "session", l.id, "err", err)
	}
	slog.Info("deliver sent mid-turn", "session", l.id, "responses", ids)
	return text
}

func responseIDs(pending []store.Pending) []string {
	ids := []string{}
	for _, p := range pending {
		ids = append(ids, p.Response.Id)
	}
	return ids
}

func deliveryMessage(pending []store.Pending) string {
	b := &strings.Builder{}
	b.WriteString("Responses from your inbox:\n")

	for _, p := range pending {
		r := p.Response
		b.WriteString("\n")

		switch p.Item.Kind {
		case api.Question:
			fmt.Fprintf(b, "Your question %s: %s\n", p.Item.Id, p.Item.Payload["question"])
		case api.SessionError:
			fmt.Fprintf(b, "Your error %s: %s\n", p.Item.Id, p.Item.Payload["error"])
		case api.RateLimited:
			fmt.Fprintf(b, "Your rate-limited turn %s: %s\n", p.Item.Id, p.Item.Payload["error"])
		}

		switch r.Type {
		case api.Option:
			if p.Item.Payload["multiple"] == true {
				fmt.Fprintf(b, "%s checked:\n", r.User)
				for _, o := range *r.Options {
					fmt.Fprintf(b, "- %s\n", o)
				}
				if r.Text != nil && strings.TrimSpace(*r.Text) != "" {
					fmt.Fprintf(b, "%s added: %s\n", r.User, *r.Text)
				}
			} else {
				fmt.Fprintf(b, "%s chose: %s\n", r.User, (*r.Options)[0])
			}
		case api.Other:
			if p.Item.Kind == api.Question {
				fmt.Fprintf(b, "%s answered (none of your options): %s\n", r.User, *r.Text)
				fmt.Fprintf(b, "If this asks you something back, answer by calling ask_human with thread_id %q.\n", p.Item.ThreadId)
			} else {
				fmt.Fprintf(b, "%s replied: %s\n", r.User, *r.Text)
			}
		}

		for _, a := range r.Attachments {
			j, err := json.Marshal(a)
			if err != nil {
				continue
			}
			fmt.Fprintf(b, "Attachment: %s\n", j)
		}
	}

	return b.String()
}