tutti/services/tuttid/service/cli/providers/agentcontext/session_get.go

package agentcontext

import (
	"context"
	"fmt"
	"strings"

	agentactivitybiz "github.com/tutti-os/tutti/packages/agent/store-sqlite"
	agentservice "github.com/tutti-os/tutti/services/tuttid/service/agent"
	cliservice "github.com/tutti-os/tutti/services/tuttid/service/cli"
)

const (
	getViewSession              = "session"
	getViewTurns                = "turns"
	getViewConversation         = "conversation"
	getViewTrace                = "trace"
	defaultConversationTurns    = 3
	defaultConversationMessages = 100
	defaultTraceMessages        = 20
)

type sessionGetResult struct {
	View           string
	Session        agentservice.Session
	Turns          []sessionGetTurnResult
	TurnSummaries  []agentactivitybiz.SessionTurnSummary
	Trace          *sessionGetTraceResult
	HasMoreTurns   bool
	ImageLocalPath imageLocalPathResolver
	WorkspaceID    string
}

type sessionGetTurnResult struct {
	Turn            agentactivitybiz.SessionTurnSummary
	Messages        []agentservice.SessionMessage
	FinalMessage    *agentservice.SessionMessage
	HasMoreMessages bool
}

type sessionGetTraceResult struct {
	Turn agentactivitybiz.SessionTurnSummary
	Page agentservice.SessionMessagesPage
}

func (p Provider) getSessionContext(ctx context.Context, workspaceID string, input getSessionInput) (any, error) {
	if err := p.requireSessions(); err != nil {
		return nil, err
	}
	view, err := validateGetSessionInput(input)
	if err != nil {
		return nil, err
	}
	if view == getViewSession {
		session, err := p.sessions.Get(ctx, workspaceID, input.SessionID)
		if err != nil {
			return nil, err
		}
		return sessionGetResult{View: view, Session: session, WorkspaceID: workspaceID}, nil
	}

	session, err := p.sessions.Get(ctx, workspaceID, input.SessionID)
	if err != nil {
		return nil, err
	}
	result := sessionGetResult{
		View:           view,
		Session:        session,
		ImageLocalPath: p.imageLocalPathResolver(ctx, workspaceID),
		WorkspaceID:    workspaceID,
	}
	if view == getViewTrace {
		trace, err := p.getSessionTrace(ctx, workspaceID, input)
		if err != nil {
			return nil, err
		}
		result.Trace = &trace
		return result, nil
	}
	if strings.TrimSpace(input.TurnID) != "" {
		turn, err := p.resolveSessionTurn(ctx, workspaceID, input.SessionID, input.TurnID)
		if err != nil {
			return nil, err
		}
		turns, err := p.getSessionConversation(ctx, workspaceID, input.SessionID, []agentactivitybiz.SessionTurnSummary{turn})
		if err != nil {
			return nil, err
		}
		result.Turns = turns
		return result, nil
	}

	page, err := p.listSessionTurns(ctx, workspaceID, input)
	if err != nil {
		return nil, err
	}
	result.HasMoreTurns = page.HasMore
	if view == getViewTurns {
		result.TurnSummaries = page.Turns
		return result, nil
	}
	turns, err := p.getSessionConversation(ctx, workspaceID, input.SessionID, page.Turns)
	if err != nil {
		return nil, err
	}
	result.Turns = turns
	return result, nil
}

func validateGetSessionInput(input getSessionInput) (string, error) {
	view := strings.TrimSpace(input.View)
	if view == "" {
		view = getViewConversation
	}
	turnID := strings.TrimSpace(input.TurnID)
	beforeTurnID := strings.TrimSpace(input.BeforeTurnID)
	switch view {
	case getViewSession:
		if input.Turns != nil || turnID != "" || beforeTurnID != "" || input.Messages != nil || input.BeforeVersion > 0 {
			return "", fmt.Errorf("%w: --view session does not accept turn or message selectors", cliservice.ErrInvalidInput)
		}
	case getViewTurns:
		if turnID != "" {
			return "", fmt.Errorf("%w: --turn-id requires --view conversation or --view trace", cliservice.ErrInvalidInput)
		}
		if input.Messages != nil || input.BeforeVersion > 0 {
			return "", fmt.Errorf("%w: --messages and --before-version require --view trace", cliservice.ErrInvalidInput)
		}
	case getViewConversation:
		if turnID != "" && (input.Turns != nil || beforeTurnID != "") {
			return "", fmt.Errorf("%w: --turn-id cannot be combined with --turns or --before-turn-id", cliservice.ErrInvalidInput)
		}
		if input.Messages != nil || input.BeforeVersion > 0 {
			return "", fmt.Errorf("%w: --messages and --before-version require --view trace", cliservice.ErrInvalidInput)
		}
	case getViewTrace:
		if turnID == "" {
			return "", fmt.Errorf("%w: --view trace requires --turn-id", cliservice.ErrInvalidInput)
		}
		if input.Turns != nil {
			return "", fmt.Errorf("%w: --turns is not supported with --view trace", cliservice.ErrInvalidInput)
		}
		if beforeTurnID != "" {
			return "", fmt.Errorf("%w: --before-turn-id is not supported with --view trace", cliservice.ErrInvalidInput)
		}
	default:
		return "", fmt.Errorf("%w: unsupported view %q", cliservice.ErrInvalidInput, view)
	}
	return view, nil
}

func (p Provider) getSessionConversation(
	ctx context.Context,
	workspaceID string,
	sessionID string,
	selected []agentactivitybiz.SessionTurnSummary,
) ([]sessionGetTurnResult, error) {
	results := make([]sessionGetTurnResult, 0, len(selected))
	for _, turn := range selected {
		page, err := p.sessions.ListMessages(ctx, workspaceID, sessionID, agentservice.ListMessagesInput{
			TurnID: turn.TurnID,
			Limit:  defaultConversationMessages,
			Order:  agentactivitybiz.MessageOrderDesc,
		})
		if err != nil {
			return nil, err
		}
		reverseSessionGetMessages(page.Messages)
		messages, final := conversationMessages(page.Messages, turn.FinalAssistantMessageID)
		if final == nil && strings.TrimSpace(turn.FinalAssistantMessageID) != "" {
			resolved, err := p.sessions.ListMessages(ctx, workspaceID, sessionID, agentservice.ListMessagesInput{
				MessageID: turn.FinalAssistantMessageID,
				Limit:     1,
				Order:     agentactivitybiz.MessageOrderDesc,
			})
			if err != nil {
				return nil, err
			}
			if len(resolved.Messages) == 1 && isConversationMessage(resolved.Messages[0]) {
				message := resolved.Messages[0]
				final = &message
			}
		}
		results = append(results, sessionGetTurnResult{
			Turn: turn, Messages: messages, FinalMessage: final, HasMoreMessages: page.HasMore,
		})
	}
	return results, nil
}

func (p Provider) listSessionTurns(ctx context.Context, workspaceID string, input getSessionInput) (agentservice.TurnPage, error) {
	limit := defaultConversationTurns
	if input.Turns != nil {
		limit = int(*input.Turns)
	}
	var before *agentactivitybiz.SessionTurnCursor
	if beforeTurnID := strings.TrimSpace(input.BeforeTurnID); beforeTurnID != "" {
		turn, err := p.resolveSessionTurn(ctx, workspaceID, input.SessionID, beforeTurnID)
		if err != nil {
			return agentservice.TurnPage{}, err
		}
		before = &agentactivitybiz.SessionTurnCursor{StartedAtUnixMS: turn.StartedAtUnixMS, TurnID: turn.TurnID}
	}
	return p.sessions.ListTurns(ctx, workspaceID, input.SessionID, agentservice.ListTurnsInput{Before: before, Limit: limit})
}

func (p Provider) getSessionTrace(
	ctx context.Context,
	workspaceID string,
	input getSessionInput,
) (sessionGetTraceResult, error) {
	turnID := strings.TrimSpace(input.TurnID)
	turn, err := p.resolveSessionTurn(ctx, workspaceID, input.SessionID, turnID)
	if err != nil {
		return sessionGetTraceResult{}, err
	}
	limit := defaultTraceMessages
	if input.Messages != nil {
		limit = int(*input.Messages)
	}
	page, err := p.sessions.ListMessages(ctx, workspaceID, input.SessionID, agentservice.ListMessagesInput{
		TurnID:        turnID,
		BeforeVersion: uint64(input.BeforeVersion),
		Limit:         limit,
		Order:         agentactivitybiz.MessageOrderDesc,
	})
	if err != nil {
		return sessionGetTraceResult{}, err
	}
	reverseSessionGetMessages(page.Messages)
	return sessionGetTraceResult{Turn: turn, Page: page}, nil
}

func (p Provider) resolveSessionTurn(ctx context.Context, workspaceID string, sessionID string, turnID string) (agentactivitybiz.SessionTurnSummary, error) {
	turnID = strings.TrimSpace(turnID)
	turn, found, err := p.sessions.GetTurn(ctx, workspaceID, sessionID, turnID)
	if err != nil {
		return agentactivitybiz.SessionTurnSummary{}, err
	}
	if !found {
		return agentactivitybiz.SessionTurnSummary{}, fmt.Errorf("%w: turn %q was not found in session", cliservice.ErrInvalidInput, turnID)
	}
	return agentactivitybiz.SessionTurnSummary{
		TurnID: turn.TurnID, Phase: turn.Phase, Outcome: turn.Outcome,
		FinalAssistantMessageID: turn.FinalAssistantMessageID,
		StartedAtUnixMS:         turn.StartedAtUnixMS, SettledAtUnixMS: turn.SettledAtUnixMS, Origin: turn.Origin,
	}, nil
}

func conversationMessages(
	messages []agentservice.SessionMessage,
	finalMessageID string,
) ([]agentservice.SessionMessage, *agentservice.SessionMessage) {
	values := make([]agentservice.SessionMessage, 0, len(messages))
	var final *agentservice.SessionMessage
	finalMessageID = strings.TrimSpace(finalMessageID)
	for _, message := range messages {
		if !isConversationMessage(message) {
			continue
		}
		if finalMessageID != "" && strings.TrimSpace(message.MessageID) == finalMessageID {
			copy := message
			final = &copy
			continue
		}
		values = append(values, message)
	}
	return values, final
}

func isConversationMessage(message agentservice.SessionMessage) bool {
	role := strings.TrimSpace(message.Role)
	if role != "user" && role != "assistant" {
		return false
	}
	if strings.TrimSpace(message.Kind) == "text" {
		return true
	}
	return role == "assistant" && message.Semantics != nil && message.Semantics.UserVisibleAssistantResponse
}

func reverseSessionGetMessages(messages []agentservice.SessionMessage) {
	for left, right := 0, len(messages)-1; left < right; left, right = left+1, right-1 {
		messages[left], messages[right] = messages[right], messages[left]
	}
}

func sessionGetJSONValue(result any) map[string]any {
	got := result.(sessionGetResult)
	value := map[string]any{
		"view":    got.View,
		"session": sessionInspectValue(got.WorkspaceID, got.Session),
	}
	addAgentSessionReference(value, got.WorkspaceID, got.Session.ID)
	switch got.View {
	case getViewTurns:
		value["turns"] = sessionGetTurnSummaryValues(got.TurnSummaries)
		value["hasMoreTurns"] = got.HasMoreTurns
	case getViewConversation:
		value["turns"] = sessionGetTurnValues(got.Turns, got.ImageLocalPath)
		value["hasMoreTurns"] = got.HasMoreTurns
	case getViewTrace:
		trace := got.Trace
		value["turn"] = sessionGetTurnValue(trace.Turn)
		value["messages"] = messageTraceValues(trace.Page.Messages, got.ImageLocalPath)
		value["latestVersion"] = trace.Page.LatestVersion
		value["hasMoreMessages"] = trace.Page.HasMore
	}
	return value
}

func sessionGetTurnSummaryValues(turns []agentactivitybiz.SessionTurnSummary) []any {
	values := make([]any, 0, len(turns))
	for _, turn := range turns {
		values = append(values, sessionGetTurnValue(turn))
	}
	return values
}

func sessionGetTurnValues(turns []sessionGetTurnResult, imageLocalPath imageLocalPathResolver) []any {
	values := make([]any, 0, len(turns))
	for _, turn := range turns {
		value := sessionGetTurnValue(turn.Turn)
		value["messages"] = messageCompactValues(turn.Messages, imageLocalPath)
		value["finalMessage"] = nil
		if turn.FinalMessage != nil {
			value["finalMessage"] = messageCompactValue(*turn.FinalMessage, imageLocalPath)
		}
		value["hasMoreMessages"] = turn.HasMoreMessages
		values = append(values, value)
	}
	return values
}

func sessionGetTurnValue(turn agentactivitybiz.SessionTurnSummary) map[string]any {
	value := map[string]any{
		"turnId": strings.TrimSpace(turn.TurnID),
		"phase":  strings.TrimSpace(turn.Phase),
	}
	if outcome := strings.TrimSpace(turn.Outcome); outcome != "" {
		value["outcome"] = outcome
	}
	if origin := strings.TrimSpace(turn.Origin); origin != "" {
		value["origin"] = origin
	}
	if turn.StartedAtUnixMS > 0 {
		value["startedAtUnixMs"] = turn.StartedAtUnixMS
	}
	if turn.SettledAtUnixMS > 0 {
		value["settledAtUnixMs"] = turn.SettledAtUnixMS
	}
	return value
}

func messageTraceValues(messages []agentservice.SessionMessage, imageLocalPath imageLocalPathResolver) []any {
	values := make([]any, 0, len(messages))
	for _, message := range messages {
		value := messageCompactValue(message, imageLocalPath)
		if len(message.Payload) > 0 {
			value["payload"] = message.Payload
		}
		if message.Semantics != nil {
			value["semantics"] = message.Semantics
		}
		if message.StartedAtUnixMS > 0 {
			value["startedAtUnixMs"] = message.StartedAtUnixMS
		}
		if message.CompletedAtUnixMS > 0 {
			value["completedAtUnixMs"] = message.CompletedAtUnixMS
		}
		values = append(values, value)
	}
	return values
}