enkly-ai/services/agent/cmd/server/main.go

// Command server runs the agent microservice: a gRPC server exposing
// api/proto/agent/v1, orchestrating vLLM-served chat + tool calls
// across the workspace, knowledge, billing and conversation services.
// Replaces internal/llm/agent.go + client.go (chat half) and the
// chat-run sections of internal/handlers/chat.go.
package main

import (
	"context"
	"log/slog"
	"net"
	"os"
	"os/signal"
	"syscall"
	"time"

	"google.golang.org/grpc"
	"google.golang.org/grpc/health"
	healthpb "google.golang.org/grpc/health/grpc_health_v1"
	"google.golang.org/grpc/reflection"

	agentv1 "enkly/gen/agent/v1"
	workspacev1 "enkly/gen/workspace/v1"
	"enkly/pkg/authsupport"
	"enkly/pkg/effort"
	"enkly/pkg/entitlements"
	"enkly/pkg/envkit"
	"enkly/pkg/vllmkit"
	"enkly/services/agent/internal/application"
	"enkly/services/agent/internal/domain"
	"enkly/services/agent/internal/infrastructure/agentrun"
	"enkly/services/agent/internal/infrastructure/grpcclients"
	"enkly/services/agent/internal/infrastructure/toolcatalog"
	"enkly/services/agent/internal/infrastructure/toolexec"
	grpcsrv "enkly/services/agent/internal/interfaces/grpc"
)

func main() {
	logger := slog.New(slog.NewTextHandler(os.Stdout, nil))

	if err := envkit.LoadDotEnv(".env"); err != nil {
		logger.Error("load .env", "err", err)
		os.Exit(1)
	}

	var (
		grpcAddr = envkit.String("ENKLY_AGENT_GRPC_ADDR", ":9006")

		workspaceAddr    = envkit.String("ENKLY_WORKSPACE_GRPC_ADDR", "workspace:9004")
		knowledgeAddr    = envkit.String("ENKLY_KNOWLEDGE_GRPC_ADDR", "knowledge:9005")
		billingAddr      = envkit.String("ENKLY_BILLING_GRPC_ADDR", "billing:9002")
		conversationAddr = envkit.String("ENKLY_CONVERSATION_GRPC_ADDR", "conversation:9003")

		// vllmChatURL / vllmEmbedURL point at the two vLLM
		// OpenAI-compatible deployments this service calls — a
		// chat-serving replica (Qwen, tool-calling + reasoning
		// enabled) and an embedding-serving replica, each its own vLLM
		// engine process since one vLLM server hosts exactly one
		// model. See k8s/11-vllm.yaml for how those are deployed (GPU
		// requests, --served-model-name, probes) and vllmkit.Client's
		// doc comment for why the two are separate fields rather than
		// one shared host the way a single local Ollama daemon used to
		// be.
		vllmChatURL  = envkit.String("ENKLY_VLLM_CHAT_URL", "http://localhost:8000/v1")
		vllmEmbedURL = envkit.String("ENKLY_VLLM_EMBED_URL", "http://localhost:8001/v1")
		chatModel    = envkit.String("ENKLY_CHAT_MODEL", vllmkit.DefaultChatModel)
		embedModel   = envkit.String("ENKLY_EMBED_MODEL", vllmkit.DefaultEmbedModel)
		// vllmAPIKey is sent as a bearer token on every request to both
		// deployments (vLLM's --api-key flag). Required once either
		// deployment is reachable from outside the cluster/mesh; leave
		// unset for a fully cluster-internal deployment with no auth in
		// front of it. See k8s/02-secrets.yaml.
		vllmAPIKey = envkit.String("ENKLY_VLLM_API_KEY", "")

		// chatThink controls vLLM/Qwen3's extended-thinking toggle
		// (vllmkit.Client.Think — "" leaves it to the chat template's
		// own default, "true"/"false" forces it, and "low"/"medium"/
		// "high" are also accepted for compatibility with pkg/effort's
		// vocabulary, though the wire protocol itself has no
		// graded-effort equivalent — see Think's doc comment).
		// chatMaxTokens/chatTemperature/chatTopP are the per-request
		// sampling knobs vLLM actually exposes over HTTP; context
		// length and GPU/tensor-parallel layout are NOT request-time
		// options under vLLM the way Ollama's num_ctx/num_gpu/
		// num_thread were — those are now fixed at server start via
		// --max-model-len/--tensor-parallel-size/
		// --gpu-memory-utilization on the vLLM deployment itself (see
		// k8s/11-vllm.yaml), so there is no client-side equivalent to
		// wire up here.
		chatThink       = envkit.String("ENKLY_CHAT_THINK", "")
		chatMaxTokens   = envkit.Int("ENKLY_CHAT_MAX_TOKENS", 0)
		chatTemperature = envkit.Float("ENKLY_CHAT_TEMPERATURE", 0)
		chatTopP        = envkit.Float("ENKLY_CHAT_TOP_P", 0)

		// Effort control / adaptive thinking (see pkg/effort and
		// agentrun.AdaptiveThinkPolicy). defaultEffort seeds every
		// Agent's fallback level when a turn carries no /effort
		// directive (see turnopts) and no WithEffort override;
		// adaptiveThinking turns on content-driven escalation above
		// that level's own floor. Both default to this deployment
		// behaving exactly as it did before effort control existed
		// (effort.Default == Medium, matching the old hardcoded
		// MaxRounds=12/DefaultThinkPolicy shape) unless explicitly
		// overridden.
		defaultEffort    = envkit.String("ENKLY_DEFAULT_EFFORT", string(effort.Default))
		adaptiveThinking = envkit.Bool("ENKLY_ADAPTIVE_THINKING", true)

		// Dynamic workflows (see dispatch_workflow / toolexec/workflow.go).
		// workflowMaxParallel is the hard operational ceiling for this
		// deployment's vLLM chat capacity; dynamicWorkflowPlans is the
		// comma-separated plan allowlist (see pkg/entitlements) —
		// "premium" today, see entitlements' package doc comment for why.
		workflowMaxParallel  = envkit.Int("ENKLY_WORKFLOW_MAX_PARALLEL", 0)
		dynamicWorkflowPlans = envkit.String("ENKLY_DYNAMIC_WORKFLOW_PLANS", "")

		// Fast mode (see /fast, Service.FastAgent). Unset
		// fastChatModel means "fast mode isn't configured on this
		// deployment" — StartTurn then just runs a /fast turn on the
		// normal Agent, exactly as if fast mode had never been
		// requested (see Service.FastAgent's doc comment). A
		// deployment that has, say, a smaller/distilled model served
		// by its own vLLM replica (or by the same replica under a
		// second --served-model-name via vLLM's multi-LoRA/multi-model
		// support) points fastChatModel at it to actually unlock the
		// latency/cost win fast mode promises; fastChatURL lets that
		// replica live on an entirely different host, which is the
		// common case (a smaller model usually means a
		// smaller/cheaper GPU shape, i.e. a different Deployment).
		fastChatModel     = envkit.String("ENKLY_FAST_CHAT_MODEL", "")
		fastChatURL       = envkit.String("ENKLY_FAST_VLLM_CHAT_URL", vllmChatURL)
		fastChatThink     = envkit.String("ENKLY_FAST_CHAT_THINK", "low")
		fastChatMaxTokens = envkit.Int("ENKLY_FAST_CHAT_MAX_TOKENS", 700)
	)

	clients, err := grpcclients.Dial(workspaceAddr, knowledgeAddr, billingAddr, conversationAddr)
	if err != nil {
		logger.Error("dial peer services", "err", err)
		os.Exit(1)
	}
	defer clients.Close()

	llmClient := vllmkit.New(vllmChatURL, vllmEmbedURL, chatModel, embedModel)
	llmClient.APIKey = vllmAPIKey
	llmClient.Think = chatThink
	llmClient.ChatOptions = vllmkit.Options{
		Temperature: chatTemperature, TopP: chatTopP, MaxTokens: chatMaxTokens,
	}

	// llmClient also backs search_office_skills' semantic index (see
	// toolexec.Executor.SkillIndex): the same vLLM-backed client
	// already used for chat satisfies skillkit.Embedder as-is.
	executor := toolexec.New(clients.Workspace, clients.Knowledge, llmClient)
	executor.WorkflowMaxParallel = workflowMaxParallel // 0 means "use toolexec's own default"

	// dispatch_subagent (and, built on the same nested Agent,
	// dispatch_workflow's parallel steps — see toolexec/workflow.go):
	// a nested Agent sharing the same vLLM-backed client but with a
	// read-only tool list (toolcatalog.SubagentDefs()) and a
	// restricted executor (toolexec.NewRestricted) that enforces that
	// list even if the model asks for anything else. Wired as a second
	// step, after executor exists, because the restricted executor has
	// to wrap executor by pointer — see toolexec.Executor.Subagent's
	// doc comment.
	executor.Subagent = agentrun.NewSubagent(llmClient, toolexec.NewRestricted(executor), toolcatalog.SubagentDefs())

	agent := agentrun.New(llmClient, executor, toolcatalog.Defs())

	// Effort control / adaptive thinking: an unrecognized
	// ENKLY_DEFAULT_EFFORT value is logged and ignored, leaving New's
	// own effort.Default (Medium) in force — a misconfigured env var
	// degrading to "the agent's long-standing default behavior" is
	// always a safe failure mode, the same stance effort.Level.Resolve
	// itself takes for an invalid value reaching it at runtime.
	if lvl, err := effort.Parse(defaultEffort); err != nil {
		logger.Warn("invalid ENKLY_DEFAULT_EFFORT, using default", "value", defaultEffort, "err", err)
	} else {
		agent.Effort = lvl
	}
	if adaptiveThinking {
		agent.AdaptiveThink = agentrun.AdaptiveThinkPolicy
	}

	// ENKLY.md project memory: if the workspace root has one, its
	// content is folded into the system prompt on every turn (see
	// agentrun.Agent.buildSystemPrompt) — the same idea as Claude Code
	// reading a CLAUDE.md. A missing file or an unreachable workspace
	// service both just mean "no project memory today"; this must never
	// fail a turn over it, so the error is swallowed here, not returned.
	const projectMemoryPath = "ENKLY.md"
	agent.MemoryLoader = func(ctx context.Context) string {
		resp, err := clients.Workspace.ReadFile(ctx, &workspacev1.ReadFileRequest{Path: projectMemoryPath})
		if err != nil || resp.GetIsError() {
			return ""
		}
		return resp.GetContent()
	}

	// Fast mode: only wired up when ENKLY_FAST_CHAT_MODEL names a
	// distinct served model — see fastChatModel's doc comment above and
	// application.Service.FastAgent's. fastAgent shares executor (and
	// therefore the same dispatch_subagent/dispatch_workflow nested
	// Agent) with the normal agent; only the ChatStreamer differs.
	var fastAgent *agentrun.Agent
	if fastChatModel != "" {
		fastClient := vllmkit.New(fastChatURL, vllmEmbedURL, fastChatModel, embedModel)
		fastClient.APIKey = vllmAPIKey
		fastClient.Think = fastChatThink
		fastClient.ChatOptions = vllmkit.Options{MaxTokens: fastChatMaxTokens}
		fastAgent = agentrun.New(fastClient, executor, toolcatalog.Defs())
		fastAgent.MemoryLoader = agent.MemoryLoader
		fastAgent.Effort = effort.Low
		logger.Info("fast mode configured", "fast_chat_model", fastChatModel, "fast_chat_url", fastChatURL)
	}

	svc := &application.Service{
		Streams:   domain.NewManager(),
		Agent:     agent,
		FastAgent: fastAgent,
		Billing:   clients.Billing,
		Conv:      clients.Conversation,
		Logger:    logger,
		// Same window/backoff as the monolith's ChatLimiter — a gentle
		// per-account throttle that "only bites a tight scripted loop".
		ChatLimiter:  authsupport.NewLimiter(time.Minute, 200*time.Millisecond, 4*time.Second),
		Entitlements: entitlements.ParsePlanList(dynamicWorkflowPlans),
	}

	lis, err := net.Listen("tcp", grpcAddr)
	if err != nil {
		logger.Error("listen", "addr", grpcAddr, "err", err)
		os.Exit(1)
	}

	grpcServer := grpc.NewServer(grpc.MaxRecvMsgSize(24<<20), grpc.MaxSendMsgSize(24<<20))
	agentv1.RegisterAgentServiceServer(grpcServer, grpcsrv.New(svc))
	healthSrv := health.NewServer()
	healthpb.RegisterHealthServer(grpcServer, healthSrv)
	healthSrv.SetServingStatus("", healthpb.HealthCheckResponse_SERVING)
	reflection.Register(grpcServer)

	go func() {
		logger.Info("agent service listening", "addr", grpcAddr, "chat_model", chatModel)
		if err := grpcServer.Serve(lis); err != nil {
			logger.Error("serve", "err", err)
			os.Exit(1)
		}
	}()

	ctx, stop := signal.NotifyContext(context.Background(), syscall.SIGINT, syscall.SIGTERM)
	defer stop()
	<-ctx.Done()
	logger.Info("shutting down")
	grpcServer.GracefulStop()
}