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"`
	Options           []string         `json:"options"`
	RecommendedOption string           `json:"recommended_option"`
	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. " +
			"Offer concrete options and recommend one. 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 what each option would concretely do. Several paragraphs is normal. Markdown is rendered."},
				"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."},
				"options": {"type": "array", "items": {"type": "string"}, "minItems": 2},
				"recommended_option": {"type": "string", "description": "One of options."},
				"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", "options", "recommended_option", "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{}
	}

	payload := map[string]any{
		"context":            args.Context,
		"question":           args.Question,
		"why_blocked":        args.WhyBlocked,
		"options":            args.Options,
		"recommended_option": args.RecommendedOption,
		"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
	m.mu.Unlock()

	m.snapshotSoon(l, api.SnapshotTriggerAskHuman)

	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")
	}
	if len(args.Options) < 2 {
		problems = append(problems, "options needs at least two concrete choices")
	}
	if !slices.Contains(args.Options, args.RecommendedOption) {
		problems = append(problems, "recommended_option must be exactly one of options")
	}
	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()

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

	waiting, err := m.store.AwaitingHuman(ctx, sessionID)
	if err != nil {
		return err
	}
	if waiting > 0 {
		return nil
	}

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

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

	_, err = m.turn(ctx, sess, task, pending[0].Response.User, deliveryMessage(pending), nil)
	if errors.Is(err, ErrConflict) {
		return nil
	}
	if err != nil {
		return err
	}

	ids := []string{}
	for _, p := range pending {
		ids = append(ids, p.Response.Id)
	}
	return m.store.MarkDelivered(ctx, 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:
			fmt.Fprintf(b, "%s chose: %s\n", r.User, *r.Option)
		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()
}