enkly-ai/services/agent/internal/application/service.go

// Package application holds the agent service's use cases — the
// transplant of internal/handlers/chat.go's Handlers.Chat (minus HTTP
// parsing/rendering) plus its runAgent background goroutine. This is a
// saga coordinator: every use case here calls out over gRPC to
// billing, conversation, workspace (via toolexec) and knowledge rather
// than touching a database of its own.
package application

import (
	"context"
	"encoding/json"
	"errors"
	"fmt"
	"log/slog"
	"strconv"
	"strings"
	"sync"
	"time"
	"unicode/utf8"

	billingv1 "enkly/gen/billing/v1"
	conversationv1 "enkly/gen/conversation/v1"
	"enkly/pkg/authsupport"
	"enkly/pkg/entitlements"
	"enkly/pkg/i18n"
	"enkly/pkg/turnopts"
	"enkly/pkg/vllmkit"
	"enkly/services/agent/internal/domain"
	"enkly/services/agent/internal/infrastructure/agentrun"
	"enkly/services/agent/internal/infrastructure/toolexec"
)

const (
	maxPromptChars  = 32_000
	historyMessages = 60 // most recent user/assistant messages replayed to the model

	// defaultConversationTitle mirrors the conversation service's own
	// DefaultConversationTitle — duplicated as a literal rather than
	// imported, since depending on another service's internal package
	// would violate the service boundary; see conversation's
	// migrations/0001_init.sql column default, which is the actual
	// source of truth both sides must agree with.
	defaultConversationTitle = "New conversation"

	maxRunDuration       = 30 * time.Minute
	sseKeepAliveInterval = 15 * time.Second
	sseCoalesceDelay     = 30 * time.Millisecond
)

// Outcome mirrors agentv1.StartTurnOutcome without importing the
// generated package into this layer — see internal/interfaces/grpc for
// the translation.
type Outcome int

const (
	OutcomeOK Outcome = iota
	OutcomeConversationNotFound
	OutcomeConversationBusy
	OutcomeTooManyStreams
	OutcomeRateLimited
	OutcomeQuotaExceeded
	OutcomeInternalError
)

type StartTurnResult struct {
	Outcome  Outcome
	Message  string
	StreamID string
}

type Service struct {
	Streams *domain.Manager
	Agent   *agentrun.Agent
	Billing billingv1.BillingServiceClient
	Conv    conversationv1.ConversationServiceClient
	Logger  *slog.Logger

	// ChatLimiter replaces Handlers.ChatLimiter: a gentle per-account
	// throttle that "only bites a tight scripted loop".
	ChatLimiter *authsupport.Limiter

	// FastAgent, when non-nil, is used instead of Agent for a turn that
	// requests fast mode (see turnopts' /fast directive) — typically
	// wired in main.go to an *agentrun.Agent sharing Executor but built
	// around a second, smaller/quantized *vllmkit.Client configured
	// for lower latency and cost per the "2.5x faster... 3x cheaper"
	// fast-mode brief. A deployment that hasn't configured a distinct
	// fast model leaves this nil, and a /fast request just runs on
	// Agent like any other turn — fast mode degrading to "no change"
	// rather than erroring is deliberate, since it's meant to be
	// available on every plan with zero configuration required.
	FastAgent *agentrun.Agent

	// Entitlements resolves which plans unlock dynamic workflows (see
	// pkg/entitlements) — built once in main.go from
	// ENKLY_DYNAMIC_WORKFLOW_PLANS and consulted on every StartTurn
	// from the plan name billing's own CheckAndIncrementUsage call
	// already returns, so this needs no gRPC call of its own. The zero
	// Catalog (an unconfigured Service, e.g. in a test) resolves every
	// plan as not entitled — see entitlements.Catalog.Resolve.
	Entitlements entitlements.Catalog
}

func cleanTitle(s string, max int) string {
	s = strings.Join(strings.Fields(s), " ")
	if utf8.RuneCountInString(s) <= max {
		return s
	}
	runes := []rune(s)
	return strings.TrimSpace(string(runes[:max-1])) + "…"
}

// titleFromPrompt replaces handlers.titleFromPrompt.
func titleFromPrompt(prompt string) string {
	first := prompt
	if i := strings.IndexByte(prompt, '\n'); i >= 0 {
		first = prompt[:i]
	}
	if t := cleanTitle(first, 60); t != "" {
		return t
	}
	return defaultConversationTitle
}

// buildHistory replaces handlers.buildHistory: earlier turns' raw tool
// results are deliberately left out — their assistant tool-call
// messages aren't persisted, so replaying the results would present
// the model with orphaned tool output; the assistant's own summary of
// what it found is already in the text.
func buildHistory(msgs []*conversationv1.Message) []vllmkit.Message {
	out := make([]vllmkit.Message, 0, len(msgs))
	for _, m := range msgs {
		role := fromProtoRole(m.GetRole())
		if (role == vllmkit.RoleUser || role == vllmkit.RoleAssistant) && strings.TrimSpace(m.GetContent()) != "" {
			out = append(out, vllmkit.Message{Role: role, Content: m.GetContent()})
		}
	}
	if len(out) > historyMessages {
		out = out[len(out)-historyMessages:]
	}
	return out
}

func fromProtoRole(r conversationv1.Role) vllmkit.Role {
	switch r {
	case conversationv1.Role_ROLE_USER:
		return vllmkit.RoleUser
	case conversationv1.Role_ROLE_ASSISTANT:
		return vllmkit.RoleAssistant
	case conversationv1.Role_ROLE_TOOL:
		return vllmkit.RoleTool
	default:
		return vllmkit.RoleSystem
	}
}

func toProtoRole(r vllmkit.Role) conversationv1.Role {
	switch r {
	case vllmkit.RoleUser:
		return conversationv1.Role_ROLE_USER
	case vllmkit.RoleAssistant:
		return conversationv1.Role_ROLE_ASSISTANT
	case vllmkit.RoleTool:
		return conversationv1.Role_ROLE_TOOL
	default:
		return conversationv1.Role_ROLE_SYSTEM
	}
}

// StartTurn replaces Handlers.Chat's synchronous prefix, launching the
// same work Handlers.runAgent did in a background goroutine.
//
// lang is the visitor's resolved UI locale (see the gateway's
// localeFrom / langFromIncomingContext in interfaces/grpc/server.go)
// — every message StartTurn itself returns is looked up in that
// locale's catalog (see pkg/i18n), and it's also threaded onto runCtx
// via agentrun.WithPreferredLanguage so the model's system prompt
// carries it as an explicit, standing language preference alongside
// its own per-message detection (see agentrun's "## Language"
// section).
func (s *Service) StartTurn(ctx context.Context, userID, convID int64, prompt string, lang i18n.Locale) StartTurnResult {
	prompt = strings.TrimSpace(prompt)
	if prompt == "" {
		return StartTurnResult{Outcome: OutcomeInternalError, Message: i18n.T(lang, "chat.type_message_first")}
	}
	if len([]rune(prompt)) > maxPromptChars {
		return StartTurnResult{Outcome: OutcomeInternalError, Message: i18n.T(lang, "chat.message_too_long", maxPromptChars)}
	}

	// Leading /effort, /fast and /system directive lines (see
	// turnopts' doc comment for exactly why this, rather than new
	// proto fields, is today's working surface for those three
	// features) are parsed off and stripped before the prompt is
	// persisted or sent to the model — a directive is a control signal
	// for this service, not part of the conversation the model or a
	// human reviewing history needs to see. prompt itself is left
	// untouched if there were no directives, so the overwhelming
	// majority of turns pay no cost here.
	opts, prompt := turnopts.Parse(prompt)
	if opts.HadDirectives {
		prompt = strings.TrimSpace(prompt)
		if prompt == "" {
			return StartTurnResult{Outcome: OutcomeInternalError, Message: i18n.T(lang, "chat.type_message_first")}
		}
	}

	if !throttle(ctx, s.ChatLimiter, "chat:"+strconv.FormatInt(userID, 10)) {
		return StartTurnResult{Outcome: OutcomeRateLimited, Message: i18n.T(lang, "chat.rate_limited")}
	}

	// Plan-based daily message cap: a lookup failure fails open (logged,
	// not blocking), matching the monolith's stance that a transient
	// billing hiccup shouldn't lock someone out of chat. planName stays
	// "" on a failed lookup, which entitlements.Catalog.Resolve treats
	// as "not entitled" for dynamic workflows — the opposite failure
	// mode from the quota check just above, and deliberately so: see
	// entitlements.Catalog.Resolve's doc comment on failing closed for
	// a gated paid feature versus failing open for message quota.
	var planName string
	if usage, err := s.Billing.CheckAndIncrementUsage(ctx, &billingv1.CheckAndIncrementUsageRequest{UserId: userID}); err != nil {
		s.Logger.Warn("check usage", "err", err, "user_id", userID)
	} else if !usage.GetAllowed() {
		return StartTurnResult{Outcome: OutcomeQuotaExceeded, Message: i18n.T(lang, "billing.quota_exceeded", usage.GetPlanName(), usage.GetLimit())}
	} else {
		planName = usage.GetPlanName()
	}

	conv, err := s.Conv.GetConversation(ctx, &conversationv1.GetConversationRequest{UserId: userID, ConversationId: convID})
	if err != nil {
		return StartTurnResult{Outcome: OutcomeConversationNotFound, Message: i18n.T(lang, "chat.conversation_gone")}
	}

	runCtx, cancel := context.WithTimeout(context.Background(), maxRunDuration)
	runCtx = agentrun.WithPreferredLanguage(runCtx, lang)
	if opts.Effort.Valid() {
		runCtx = agentrun.WithEffort(runCtx, opts.Effort)
	}
	if len(opts.SystemEntries) > 0 {
		runCtx = agentrun.WithSystemEntries(runCtx, toAgentrunSystemEntries(opts.SystemEntries))
	}
	runCtx = toolexec.WithEntitlements(runCtx, s.Entitlements.Resolve(planName))

	// Fast mode (see /fast and Service.FastAgent's doc comment) picks
	// which Agent actually runs this turn — everything else about the
	// run (context, callbacks, persistence) is identical either way.
	runAgent := s.Agent
	if opts.Fast && s.FastAgent != nil {
		runAgent = s.FastAgent
	}

	stream, err := s.Streams.Create(userID, conv.GetId(), cancel)
	if err != nil {
		cancel()
		switch {
		case errors.Is(err, domain.ErrConversationBusy):
			return StartTurnResult{Outcome: OutcomeConversationBusy, Message: i18n.T(lang, "chat.conversation_busy")}
		case errors.Is(err, domain.ErrTooManyStreams):
			return StartTurnResult{Outcome: OutcomeTooManyStreams, Message: i18n.T(lang, "chat.too_many_streams")}
		default:
			s.Logger.Error("create stream", "err", err)
			return StartTurnResult{Outcome: OutcomeInternalError, Message: i18n.T(lang, "errors.generic")}
		}
	}

	if _, err := s.Conv.AddMessage(ctx, &conversationv1.AddMessageRequest{
		ConversationId: conv.GetId(), Role: conversationv1.Role_ROLE_USER, Content: prompt,
	}); err != nil {
		cancel()
		stream.Finish("")
		s.Streams.ScheduleRemoval(stream)
		s.Logger.Error("save user message", "err", err)
		return StartTurnResult{Outcome: OutcomeInternalError, Message: i18n.T(lang, "chat.save_failed")}
	}
	if _, err := s.Conv.AutoTitleConversation(ctx, &conversationv1.AutoTitleConversationRequest{
		UserId: userID, ConversationId: conv.GetId(), Title: titleFromPrompt(prompt),
	}); err != nil {
		s.Logger.Warn("auto-title conversation", "err", err)
	}

	msgsResp, err := s.Conv.ListMessages(ctx, &conversationv1.ListMessagesRequest{ConversationId: conv.GetId()})
	if err != nil {
		cancel()
		stream.Finish("")
		s.Streams.ScheduleRemoval(stream)
		s.Logger.Error("load history", "err", err)
		return StartTurnResult{Outcome: OutcomeInternalError, Message: i18n.T(lang, "chat.load_failed")}
	}

	go s.runTurn(runCtx, cancel, stream, conv.GetId(), buildHistory(msgsResp.GetMessages()), runAgent)

	return StartTurnResult{Outcome: OutcomeOK, StreamID: stream.ID}
}

// toAgentrunSystemEntries translates turnopts.SystemEntry (a
// transport-agnostic shape pkg/turnopts owns so it doesn't need to
// import the agent service's internal agentrun package) to
// agentrun.SystemEntry — the one call site that actually needs both
// types in scope.
func toAgentrunSystemEntries(in []turnopts.SystemEntry) []agentrun.SystemEntry {
	out := make([]agentrun.SystemEntry, 0, len(in))
	for _, e := range in {
		var kind agentrun.SystemEntryKind
		switch e.Kind {
		case turnopts.SystemEntryPermission:
			kind = agentrun.SystemEntryPermission
		case turnopts.SystemEntryBudget:
			kind = agentrun.SystemEntryBudget
		default:
			kind = agentrun.SystemEntryInstruction
		}
		out = append(out, agentrun.SystemEntry{Kind: kind, Content: e.Content})
	}
	return out
}

// recordedCall mirrors handlers.recordedCall — used to match a
// persisted tool message back to its structured call/result for the
// tool_calls audit log, since callbacks can fire out of order when
// tools run concurrently.
type recordedCall struct {
	call   vllmkit.ToolCallRequest
	result vllmkit.ToolResult
	used   bool
}

// runTurn replaces Handlers.runAgent. runAgent is the *agentrun.Agent
// StartTurn selected for this turn — s.Agent normally, or s.FastAgent
// when fast mode was requested and configured (see StartTurn) — rather
// than a hardcoded s.Agent, so fast mode can swap in a differently
// configured Agent with zero other change to how a turn runs.
func (s *Service) runTurn(ctx context.Context, cancel context.CancelFunc, stream *domain.Stream, convID int64, history []vllmkit.Message, runAgent *agentrun.Agent) {
	defer cancel()
	defer s.Streams.ScheduleRemoval(stream)

	defer func() {
		if rec := recover(); rec != nil {
			s.Logger.Error("panic in agent run", "panic", rec, "conversation_id", convID)
			stream.Finish("⚠️ Something went wrong while running the agent.")
		}
	}()

	var (
		mu    sync.Mutex
		calls []*recordedCall
	)
	cb := agentrun.Callbacks{
		OnToken:         stream.SetToken,
		OnThinking:      stream.SetThinking,
		OnThinkingRound: stream.AddThinkingRound,
		OnToolCall: func(call vllmkit.ToolCallRequest, result vllmkit.ToolResult) {
			mu.Lock()
			calls = append(calls, &recordedCall{call: call, result: result})
			mu.Unlock()
			stream.AddTool(domain.ToolCallResult{
				Name: result.Name, ArgumentsJSON: argsJSON(call.Arguments), Result: result.Content, IsError: result.IsError,
			})
		},
	}

	final, toPersist, runErr := runAgent.Run(ctx, history, cb)

	pctx, pcancel := context.WithTimeout(context.Background(), 30*time.Second)
	defer pcancel()

	note := ""
	switch {
	case runErr == nil:
	case agentrun.IsCancelled(runErr) && errors.Is(ctx.Err(), context.DeadlineExceeded):
		note = fmt.Sprintf("⏱ Stopped: this reply hit the %s time limit.", maxRunDuration)
	case agentrun.IsCancelled(runErr):
		note = "⏹ Stopped."
	case vllmkit.IsUpstreamUnavailable(runErr):
		// The vLLM chat deployment is unreachable (Pending/restarting/
		// mid-rollout) — a raw "dial tcp ...: connection refused"
		// tells the end user nothing actionable, so show a short
		// status message instead and keep the real error in the logs
		// for whoever runs k8s/scripts/gpu-preflight.sh.
		s.Logger.Error("agent run failed: vllm upstream unavailable", "err", runErr, "conversation_id", convID)
		note = "⚠️ The AI model service is starting up or temporarily unavailable. Please try again in a minute."
	default:
		s.Logger.Error("agent run failed", "err", runErr, "conversation_id", convID)
		note = "⚠️ " + runErr.Error()
	}
	if note != "" {
		toPersist = append(toPersist, vllmkit.Message{Role: vllmkit.RoleAssistant, Content: note})
	}

	for _, m := range toPersist {
		resp, err := s.Conv.AddMessage(pctx, &conversationv1.AddMessageRequest{
			ConversationId: convID, Role: toProtoRole(m.Role), Content: m.Content, ToolName: m.ToolName,
		})
		if err != nil {
			s.Logger.Error("persist message", "err", err, "conversation_id", convID)
			continue
		}
		if m.Role != vllmkit.RoleTool {
			continue
		}
		mu.Lock()
		for _, rc := range calls {
			if !rc.used && rc.call.Name == m.ToolName && rc.result.Content == m.Content {
				rc.used = true
				if _, err := s.Conv.AddToolCall(pctx, &conversationv1.AddToolCallRequest{
					MessageId: resp.GetId(), ToolName: rc.call.Name, ArgumentsJson: argsJSON(rc.call.Arguments),
					Result: rc.result.Content, IsError: rc.result.IsError,
				}); err != nil {
					s.Logger.Warn("persist tool call", "err", err)
				}
				break
			}
		}
		mu.Unlock()
	}

	display := strings.TrimSpace(final)
	if note != "" {
		if display != "" {
			display += "\n\n"
		}
		display += note
	}
	if display == "" {
		display = "*The model returned an empty response.*"
	}

	var toolCalls []domain.ToolCallResult
	mu.Lock()
	for _, rc := range calls {
		toolCalls = append(toolCalls, domain.ToolCallResult{
			Name: rc.call.Name, ArgumentsJSON: argsJSON(rc.call.Arguments), Result: rc.result.Content, IsError: rc.result.IsError,
		})
	}
	mu.Unlock()

	stream.FinishWithToolCalls(display, toolCalls)
}

// StopTurn replaces Handlers.StopStream. Returns false if streamID
// doesn't belong to userID (or doesn't exist), exactly like the
// monolith's http.NotFound check.
func (s *Service) StopTurn(userID int64, streamID string) bool {
	stream := s.Streams.Get(streamID)
	if stream == nil || stream.UserID != userID {
		return false
	}
	stream.Cancel()
	return true
}

func argsJSON(args map[string]any) string {
	// Matches llm.ArgsJSON's fallback-to-"{}" behavior exactly.
	b, err := json.Marshal(args)
	if err != nil {
		return "{}"
	}
	return string(b)
}