byre/internal/agents/inbox.go

package agents

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

	"g.fc.run/ai/byre/internal/api"
	"g.fc.run/ai/byre/internal/mcp"
	"g.fc.run/ai/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"`
	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: " +
			"`- ( ) 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. " +
			"Everything specific to one choice goes inside that choice, never in the text above the list: " +
			"its first line is a short name, and the lines indented beneath it say what it would concretely do, " +
			"its example, code, commands and trade-offs, so the human never has to match a paragraph above to an option below. " +
			"Text above the list covers only what all the choices share. " +
			"To choose between images (icons, screenshots, mockups), give each choice exactly one indented line, its image " +
			"(`![alt](path)`, the path relative to your worktree or absolute inside it; byre serves no image from outside the worktree, " +
			"so an image in a scratch or temp directory shows as \"image not available\": copy it into the worktree first), and nothing else: the choices then show side by side as a grid " +
			"so the human can compare them at a glance. " +
			"End your turn after asking; the answer arrives as your next message. " +
			"An answer of \"Use your judgment for this decision and continue within the existing task instructions.\" lets you settle that one decision yourself; " +
			"it never grants an authorization the owner must give explicitly, such as to deploy, push or commit: ask for that on its own. " +
			"One call is one decision, answered by picking from its own choices. When you have several decisions, such as a design with several choices, " +
			"never bundle them into one call that asks for approval of them all: call this once per decision, each with its own context and choices; 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 shared by all the choices, and then the choices. Markdown is rendered. Each choice is a list item, at least two, all of one kind: '- ( ) Merge now' / '- (*) 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 first line, after the marker, is a short name that comes back verbatim as the answer. Everything about that one choice goes in lines indented two spaces beneath it, not in a paragraph above the list: what it would concretely do, its worked example, any code block, command or sub-list, and its trade-offs. For example:\n\n- ( ) Retry in the client\n  Wraps ` + "`fetch` in `retry(3)`" + `; a 500 shows the spinner up to 6s longer.\n  ` + "```" + `ts\n  await retry(3, () => api.get(path));\n  ` + "```" + `\n- (*) Queue on the server\n  The request returns at once; the write lands within 1s.\n\nPut 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"]
				}},
				"thread_id": {"type": "string"}
			},
			"required": ["context", "question", "why_blocked"]
		}`),
		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
	}

	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, api.InboxItem{
		Id:             id,
		ProjectId:      l.projectID,
		TaskId:         l.taskID,
		AgentSessionId: l.id,
		ThreadId:       thread,
		Payload:        store.ToPayload(payload),
		State:          api.Open,
		CreatedAt:      time.Now().UTC(),
		Responses:      []api.InboxResponse{},
	})
	if err != nil {
		return "", err
	}

	m.mu.Lock()
	l.asked = true
	m.mu.Unlock()

	return fmt.Sprintf("Question filed in the human's inbox. question_id: %s. End your turn now. The answer will arrive as your next message. "+
		"If the human answers it some other way, such as in a chat message, withdraw it with cancel_question and question_id %q.", id, id), 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 {
		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`")
	}
	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) recordError(l *live, msg string) {
	err := m.emit(context.Background(), l, "byre", "session.error", map[string]any{"error": msg})
	if err != nil {
		slog.Error("record session error", "session", l.id, "err", err)
	}
}

func (m *Manager) recordRateLimited(l *live, info map[string]any, failure string) {
	payload := map[string]any{"vendor": string(l.vendor), "rate_limit": info, "error": failure}
	err := m.emit(context.Background(), l, "byre", "session.rate_limited", payload)
	if err != nil {
		slog.Error("record rate limit", "session", l.id, "err", err)
	}
}

func (m *Manager) Deliver(ctx context.Context, sessionID string) error {
	unlock := m.lockDelivery(sessionID)
	defer unlock()

	sess, err := m.store.GetSession(ctx, sessionID)
	if err != nil {
		return err
	}
	if sess.State == api.AgentSessionStateIdle {
		sent, err := m.deliverQueued(ctx, sess)
		if err != nil || sent {
			return err
		}
	}

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

	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: question waiting", "session", sessionID, "waiting", waiting, "responses", ids)
		return nil
	}

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

	target, err := m.sessionTarget(ctx, sess, task)
	if err != nil {
		return err
	}
	snapshotN := m.snapshotForAnswer(ctx, target, false)

	images := []string{}
	for _, p := range pending {
		images = append(images, p.Response.Images...)
	}
	after, err := m.turn(ctx, sess, task, pending[0].Response.User, deliveryMessage(pending), images)
	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, "snapshot", snapshotN)
	return m.store.MarkDelivered(ctx, ids, snapshotN)
}

type midTurn struct {
	text       string
	taskInputs []store.TaskInput
	leadInputs []store.LeadInput
	responses  []string
}

func (m *Manager) deliverMidTurn(ctx context.Context, l *live) midTurn {
	if l.owner != "" {
		inputs := m.takeLeadQueuedMidTurn(ctx, l)
		if len(inputs) == 0 {
			return midTurn{}
		}
		msgs := []queued{}
		for _, in := range inputs {
			msgs = append(msgs, queued{actor: in.Actor, text: in.Text})
		}
		return midTurn{text: m.recordQueued(ctx, l, msgs), leadInputs: inputs}
	}

	unlock := m.lockDelivery(l.id)
	defer unlock()

	discard, err := m.store.GetDiscard(ctx, l.taskID)
	if err != nil {
		slog.Error("read discard state", "task", l.taskID, "err", err)
		return midTurn{}
	}
	if discard.State != store.DiscardNone {
		return midTurn{}
	}

	mt := midTurn{}
	parts := []string{}
	mt.taskInputs = m.takeQueuedMidTurn(ctx, l)
	if len(mt.taskInputs) > 0 {
		msgs := []queued{}
		for _, in := range mt.taskInputs {
			msgs = append(msgs, queued{actor: in.Actor, text: in.Text})
		}
		parts = append(parts, m.recordQueued(ctx, l, msgs))
	}
	responses, ids := m.responsesMidTurn(ctx, l)
	if responses != "" {
		parts = append(parts, responses)
		mt.responses = ids
	}
	mt.text = strings.Join(parts, "\n\n")
	return mt
}

func (m *Manager) restoreMidTurn(ctx context.Context, l *live, mt midTurn) {
	if len(mt.taskInputs) > 0 {
		err := m.store.RestoreTaskInputs(ctx, mt.taskInputs)
		if err != nil {
			slog.Error("restore queued messages", "session", l.id, "err", err)
		}
	}
	if len(mt.leadInputs) > 0 {
		err := m.store.RestoreLeadInputs(ctx, mt.leadInputs)
		if err != nil {
			slog.Error("restore held lead messages", "session", l.id, "err", err)
		}
	}
	if len(mt.responses) > 0 {
		err := m.store.MarkUndelivered(ctx, mt.responses)
		if err != nil {
			slog.Error("mark responses undelivered", "session", l.id, "err", err, "responses", mt.responses)
		}
	}
	slog.Info("mid-turn delivery restored", "session", l.id, "inputs", len(mt.taskInputs)+len(mt.leadInputs), "responses", mt.responses)
}

func (m *Manager) responsesMidTurn(ctx context.Context, l *live) (string, []string) {
	pending, err := m.store.Undelivered(ctx, l.id)
	if err != nil {
		slog.Error("list undelivered responses", "session", l.id, "err", err)
		return "", nil
	}
	if len(pending) == 0 {
		return "", nil
	}
	ids := responseIDs(pending)
	for _, p := range pending {
		if len(p.Response.Images) > 0 {
			slog.Info("deliver held mid-turn: response has images", "session", l.id, "responses", ids)
			return "", nil
		}
	}

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

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

	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, "snapshot", snapshotN)
	return text, ids
}

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")

		fmt.Fprintf(b, "Your question %s: %s\n", p.Item.Id, p.Item.Payload["question"])

		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:
			text := ""
			if r.Text != nil {
				text = *r.Text
			}
			fmt.Fprintf(b, "%s answered (none of your options): %s\n", r.User, text)
			fmt.Fprintf(b, "If this asks you something back, answer by calling ask_human with thread_id %q.\n", p.Item.ThreadId)
		}

		if len(r.Images) > 0 {
			fmt.Fprintf(b, "%s attached %d image(s) to this response; they come with this message.\n", r.User, len(r.Images))
		}

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

	return b.String()
}