enkly-ai/services/agent/internal/interfaces/grpc/server.go

// Package grpc adapts application.Service onto the generated
// agentv1.AgentServiceServer interface. StreamTurn's polling/diffing/
// coalescing loop lives here rather than in the application layer
// because it's intrinsically about the generated grpc.ServerStream
// type — see its doc comment for how it maps to the monolith's
// Handlers.StreamChat.
package grpc

import (
	"context"
	"time"

	"google.golang.org/grpc/metadata"

	agentv1 "enkly/gen/agent/v1"
	"enkly/pkg/grpcerr"
	"enkly/pkg/i18n"
	"enkly/services/agent/internal/application"
	"enkly/services/agent/internal/domain"
)

const (
	sseKeepAliveInterval = 15 * time.Second
	sseCoalesceDelay     = 30 * time.Millisecond

	// langMetadataKey mirrors the gateway's handlers.langMetadataKey —
	// the gRPC metadata key it sets on StartTurn to tell this service
	// which UI locale the visitor has chosen. Kept as a plain string
	// constant rather than a shared pkg symbol because it's the one
	// and only coupling point between the two services' transport
	// layers, and a shared constant for a single string isn't worth a
	// new shared pkg import on either side.
	langMetadataKey = "x-enkly-lang"
)

// langFromIncomingContext reads the visitor's chosen locale off ctx's
// incoming gRPC metadata (set by the gateway — see langMetadataKey's
// doc comment), defaulting to i18n.Default when the call carries none
// (any caller other than the gateway, or an older gateway build).
func langFromIncomingContext(ctx context.Context) i18n.Locale {
	md, ok := metadata.FromIncomingContext(ctx)
	if !ok {
		return i18n.Default
	}
	values := md.Get(langMetadataKey)
	if len(values) == 0 {
		return i18n.Default
	}
	return i18n.ParseLocale(values[0])
}

type Server struct {
	agentv1.UnimplementedAgentServiceServer
	Svc *application.Service
}

func New(svc *application.Service) *Server { return &Server{Svc: svc} }

func toProtoOutcome(o application.Outcome) agentv1.StartTurnOutcome {
	switch o {
	case application.OutcomeOK:
		return agentv1.StartTurnOutcome_START_TURN_OUTCOME_OK
	case application.OutcomeConversationNotFound:
		return agentv1.StartTurnOutcome_START_TURN_OUTCOME_CONVERSATION_NOT_FOUND
	case application.OutcomeConversationBusy:
		return agentv1.StartTurnOutcome_START_TURN_OUTCOME_CONVERSATION_BUSY
	case application.OutcomeTooManyStreams:
		return agentv1.StartTurnOutcome_START_TURN_OUTCOME_TOO_MANY_STREAMS
	case application.OutcomeRateLimited:
		return agentv1.StartTurnOutcome_START_TURN_OUTCOME_RATE_LIMITED
	case application.OutcomeQuotaExceeded:
		return agentv1.StartTurnOutcome_START_TURN_OUTCOME_QUOTA_EXCEEDED
	default:
		return agentv1.StartTurnOutcome_START_TURN_OUTCOME_INTERNAL_ERROR
	}
}

func (s *Server) StartTurn(ctx context.Context, req *agentv1.StartTurnRequest) (*agentv1.StartTurnResponse, error) {
	lang := langFromIncomingContext(ctx)
	res := s.Svc.StartTurn(ctx, req.GetUserId(), req.GetConversationId(), req.GetPrompt(), lang)
	return &agentv1.StartTurnResponse{Outcome: toProtoOutcome(res.Outcome), Message: res.Message, StreamId: res.StreamID}, nil
}

func toProtoToolCall(tc domain.ToolCallResult) *agentv1.ToolCall {
	return &agentv1.ToolCall{Name: tc.Name, ArgumentsJson: tc.ArgumentsJSON, Result: tc.Result, IsError: tc.IsError}
}

func toProtoThinkingRound(tr domain.ThinkingRound) *agentv1.ThinkingRound {
	return &agentv1.ThinkingRound{Round: int32(tr.Round), Thinking: tr.Thinking}
}

// StreamTurn replaces Handlers.StreamChat, translated from SSE to
// server-streaming gRPC: same snapshot/diff/coalesce loop, same
// Last-Event-ID-style resume (resume_from_tool_index), same keepalive
// interval — just agentv1.TurnEvent messages instead of SSE frames. The
// gateway's own SSE handler reads this stream and does the reverse
// translation back to SSE for the browser, plus all HTML rendering
// (this service deals in plain text only).
func (s *Server) StreamTurn(req *agentv1.StreamTurnRequest, stream agentv1.AgentService_StreamTurnServer) error {
	st := s.Svc.Streams.Get(req.GetStreamId())
	if st == nil || st.UserID != req.GetUserId() {
		return grpcerr.NotFound("stream")
	}

	sentTools := int(req.GetResumeFromToolIndex())
	// sentThinkRounds has no resume-from-request counterpart the way
	// sentTools does (StreamTurnRequest carries no thinking-round
	// cursor) — a reconnect just resends every round from the start,
	// same as lastThinking resetting to "" below and re-sending the
	// full cumulative trace. Cheap and harmless: a run has at most
	// MaxRounds (12) of these, each a short text, never the
	// unboundedly-large payloads tool results can be.
	sentThinkRounds := 0
	var lastToken, lastThinking string

	ctx := stream.Context()
	for {
		snap, changed := st.Snapshot()

		for i := sentTools; i < len(snap.Tools); i++ {
			if err := stream.Send(&agentv1.TurnEvent{Event: &agentv1.TurnEvent_ToolCall{ToolCall: toProtoToolCall(snap.Tools[i])}}); err != nil {
				return err
			}
		}
		if len(snap.Tools) > sentTools {
			sentTools = len(snap.Tools)
		}
		for i := sentThinkRounds; i < len(snap.ThinkingRounds); i++ {
			if err := stream.Send(&agentv1.TurnEvent{Event: &agentv1.TurnEvent_ThinkingRound{ThinkingRound: toProtoThinkingRound(snap.ThinkingRounds[i])}}); err != nil {
				return err
			}
		}
		if len(snap.ThinkingRounds) > sentThinkRounds {
			sentThinkRounds = len(snap.ThinkingRounds)
		}
		if snap.ThinkText != lastThinking {
			lastThinking = snap.ThinkText
			if err := stream.Send(&agentv1.TurnEvent{Event: &agentv1.TurnEvent_ThinkingDelta{ThinkingDelta: lastThinking}}); err != nil {
				return err
			}
		}
		if snap.Done {
			done := &agentv1.TurnDone{FinalText: snap.FinalText}
			for _, tc := range snap.FinalToolCalls {
				done.ToolCalls = append(done.ToolCalls, toProtoToolCall(tc))
			}
			return stream.Send(&agentv1.TurnEvent{Event: &agentv1.TurnEvent_Done{Done: done}})
		}
		if snap.TokenText != lastToken {
			lastToken = snap.TokenText
			if err := stream.Send(&agentv1.TurnEvent{Event: &agentv1.TurnEvent_TokenDelta{TokenDelta: lastToken}}); err != nil {
				return err
			}
		}

		select {
		case <-changed:
			// Let a burst of tokens accumulate so we send one event, not fifty.
			select {
			case <-time.After(sseCoalesceDelay):
			case <-ctx.Done():
				return nil
			}
		case <-time.After(sseKeepAliveInterval):
			// Server-streaming gRPC has its own keepalive at the
			// transport level (HTTP/2 PINGs); nothing to send here.
			// The timer's only job is to re-loop and re-check ctx.Done.
		case <-ctx.Done():
			return nil
		}
	}
}

func (s *Server) StopTurn(ctx context.Context, req *agentv1.StopTurnRequest) (*agentv1.StopTurnResponse, error) {
	if !s.Svc.StopTurn(req.GetUserId(), req.GetStreamId()) {
		return nil, grpcerr.NotFound("stream")
	}
	return &agentv1.StopTurnResponse{}, nil
}