okami.mom/albedo

git clone https://git.nkp.pet/okami.mom/albedo

nekomimi.petfeat: multiple images, collapsed pastes, and an attachment strip in the prompt6e828112026-10-04

package app

import (
	"context"
	"errors"
	"fmt"
	"slices"
	"time"

	"albedo/cli/internal/daemon"
)

type PromptOptions struct {
	Prompt, SessionID, Workspace, Model string
	Timeout                             time.Duration
}
type PromptResult struct{ SessionID, Answer string }

// RunPrompt waits for the reply to this invocation's turn. Zero timeout is unlimited.
func (s *Service) RunPrompt(ctx context.Context, options PromptOptions) (PromptResult, error) {
	prompt, sessionID, cwd, model, timeout := options.Prompt, options.SessionID, options.Workspace, options.Model, options.Timeout
	if timeout > 0 {
		var cancel context.CancelFunc
		ctx, cancel = context.WithTimeout(ctx, timeout)
		defer cancel()
	}
	conn, err := s.Connect(ctx)
	if err != nil {
		return PromptResult{}, err
	}
	var choice modelChoice
	if model != "" {
		profiles, profileErr := daemon.ProviderProfiles(ctx, conn)
		if profileErr != nil {
			return PromptResult{}, profileErr
		}
		if choice, err = chooseModel(configuredModels(ctx, conn, profiles), profiles.Active, model); err != nil {
			return PromptResult{}, err
		}
	}
	if sessionID != "" {
		if sessionID, err = resolveSession(ctx, conn, sessionID); err != nil {
			return PromptResult{}, err
		}
	}
	if sessionID == "" {
		create := daemon.CreateSessionRequest{Workspace: cwd}
		if model != "" {
			create.Provider, create.Model = choice.provider, choice.model
		}
		session, createErr := daemon.CreateSession(ctx, conn, create)
		if createErr != nil {
			return PromptResult{}, createErr
		}
		sessionID = session.ID
	} else if model != "" {
		if switchErr := switchModel(ctx, conn, sessionID, choice); switchErr != nil {
			return PromptResult{}, switchErr
		}
	}

	client := daemon.NewChatClient(conn, sessionID)
	answer, err := awaitReply(ctx, client, prompt)
	if err == nil {
		return PromptResult{SessionID: sessionID, Answer: answer}, nil
	}
	switch {
	case errors.Is(ctx.Err(), context.DeadlineExceeded):
		return PromptResult{}, fmt.Errorf("timed out after %s; %w; the session is %s", timeout, err, sessionID)
	case ctx.Err() != nil:
		return PromptResult{}, fmt.Errorf("interrupted; %w; the session is %s", err, sessionID)
	}
	if _, uncertain := errors.AsType[*daemon.UncertainOutcomeError](err); uncertain {
		return PromptResult{}, fmt.Errorf("%w; the session is %s", err, sessionID)
	}
	return PromptResult{}, err
}

// awaitReply follows the logical turn containing this submission. The actor
// publishes membership before any worker events, and completion after them.
func awaitReply(ctx context.Context, client *daemon.ChatClient, prompt string) (answer string, failure error) {
	handle, err := client.PrepareTurn(prompt, nil, nil, false)
	if err != nil {
		return "", err
	}
	submissionID := handle.ID()
	defer func() {
		if failure != nil && ctx.Err() != nil {
			if _, uncertain := errors.AsType[*daemon.UncertainOutcomeError](failure); !uncertain {
				uncertain := &daemon.UncertainOutcomeError{Operation: "await turn", Handle: handle, Cause: failure}
				failure = fmt.Errorf("%w; %v", uncertain, failure)
			}
		}
	}()
	streamCtx, cancel := context.WithCancel(ctx)
	defer cancel()
	submitted := false
	defer func() {
		if failure == nil || !submitted {
			return
		}
		cleanupCtx, stop := context.WithTimeout(context.Background(), 5*time.Second)
		defer stop()
		outcome, err := client.CancelSubmission(cleanupCtx, submissionID)
		if err != nil {
			failure = fmt.Errorf("%w; cancellation unconfirmed: %v", failure, err)
			return
		}
		switch outcome {
		case "cancelled_queued":
			failure = fmt.Errorf("%w; queued submission cancelled", failure)
		case "interrupt_requested":
			failure = fmt.Errorf("%w; interruption requested for the turn", failure)
		case "shared_running":
			failure = fmt.Errorf("%w; the shared turn is still running", failure)
		case "not_pending":
			failure = fmt.Errorf("%w; submission is no longer pending", failure)
		}
	}()
	ready := make(chan struct{})
	events := make(chan daemon.StreamEvent, 64)
	streamDone := make(chan error, 1)
	go func() {
		streamDone <- client.Stream(streamCtx, 0, func(event daemon.StreamEvent) error {
			if event.Type == daemon.EventReset {
				select {
				case <-ready:
				default:
					close(ready)
				}
			}
			if event.Replayed {
				return nil
			}
			select {
			case events <- event:
				return nil
			case <-streamCtx.Done():
				return streamCtx.Err()
			}
		})
	}()
	select {
	case <-ready:
	case err := <-streamDone:
		return "", fmt.Errorf("could not follow the session: %w", err)
	case <-ctx.Done():
		return "", ctx.Err()
	}
	// A failed request can already have reached the actor. Cleanup always uses
	// its identity, so it cannot interrupt a different client's turn.
	submitted = true
	if _, err := client.SubmitOperation(ctx, handle); err != nil {
		return "", err
	}
	turnID := ""
	// Discard the initial reset already acknowledged before submission.
	initialReset := true
	for {
		select {
		case <-ctx.Done():
			return "", ctx.Err()
		case err := <-streamDone:
			return "", &daemon.UncertainOutcomeError{Operation: "await turn", Handle: handle, Cause: fmt.Errorf("lost the session stream before the turn finished: %v", err)}
		case event := <-events:
			if event.Type == daemon.EventReset && initialReset {
				initialReset = false
				continue
			}
			if event.Type == daemon.EventReset {
				return "", &daemon.UncertainOutcomeError{Operation: "await turn", Handle: handle, Cause: errors.New("session stream reset before the submission completed")}
			}
			if event.Type == daemon.EventTurnMembership && slices.Contains(event.SubmissionIDs, submissionID) {
				turnID = event.TurnID
			}
			if turnID == "" || event.TurnID != "" && event.TurnID != turnID {
				continue
			}
			switch event.Type {
			case daemon.EventMessage:
				answer, failure = event.Text, nil
			case daemon.EventError:
				failure = errors.New(event.Text)
			case daemon.EventInterrupted:
				failure = errors.New("the turn was interrupted")
			case daemon.EventRetry:
				failure = nil
			case daemon.EventTurnCompleted:
				if event.TurnID != turnID {
					continue
				}
				if event.Source == "interrupted" {
					return "", errors.New("the turn was interrupted")
				}
				if event.Source == "failed" && failure == nil {
					failure = errors.New("the turn failed")
				}
				if failure != nil {
					return "", failure
				}
				if answer == "" {
					return "", errors.New("the turn ended without an assistant reply")
				}
				return answer, nil
			}
		}
	}
}