707 lines
26 KiB
Go
707 lines
26 KiB
Go
package main
|
|
|
|
import (
|
|
"context"
|
|
"encoding/json"
|
|
"errors"
|
|
"fmt"
|
|
"log/slog"
|
|
"net"
|
|
"net/http"
|
|
"os"
|
|
"os/signal"
|
|
"path/filepath"
|
|
"strings"
|
|
"syscall"
|
|
"time"
|
|
|
|
agentv1 "git.ipao.vip/rogee/go-sip/gen/agent/v1"
|
|
"git.ipao.vip/rogee/go-sip/internal/agent"
|
|
"git.ipao.vip/rogee/go-sip/internal/ai"
|
|
"git.ipao.vip/rogee/go-sip/internal/callflow"
|
|
"git.ipao.vip/rogee/go-sip/internal/calllog"
|
|
"git.ipao.vip/rogee/go-sip/internal/callruntime"
|
|
"git.ipao.vip/rogee/go-sip/internal/callwindow"
|
|
"git.ipao.vip/rogee/go-sip/internal/config"
|
|
"git.ipao.vip/rogee/go-sip/internal/contract"
|
|
"git.ipao.vip/rogee/go-sip/internal/control"
|
|
"git.ipao.vip/rogee/go-sip/internal/dispatcher"
|
|
"git.ipao.vip/rogee/go-sip/internal/health"
|
|
"git.ipao.vip/rogee/go-sip/internal/mq"
|
|
ossclient "git.ipao.vip/rogee/go-sip/internal/oss"
|
|
"git.ipao.vip/rogee/go-sip/internal/rpc"
|
|
"git.ipao.vip/rogee/go-sip/internal/store"
|
|
"github.com/spf13/cobra"
|
|
"google.golang.org/grpc"
|
|
"google.golang.org/grpc/credentials"
|
|
)
|
|
|
|
func main() {
|
|
if err := newRootCommand().Execute(); err != nil {
|
|
slog.Error("command failed", "error", err)
|
|
os.Exit(1)
|
|
}
|
|
}
|
|
|
|
func newRootCommand() *cobra.Command {
|
|
root := &cobra.Command{
|
|
Use: "sip-go-agent",
|
|
Short: "SIP Go Agent and Dispatcher",
|
|
SilenceUsage: true,
|
|
SilenceErrors: true,
|
|
}
|
|
root.AddCommand(newAgentCommand(), newDispatcherCommand())
|
|
return root
|
|
}
|
|
|
|
func newAgentCommand() *cobra.Command {
|
|
cfg := config.FromEnv()
|
|
var realCall bool
|
|
cmd := &cobra.Command{
|
|
Use: "agent",
|
|
Short: "run the file-backed Agent process",
|
|
RunE: func(_ *cobra.Command, _ []string) error {
|
|
if realCall {
|
|
return runCallOnce(cfg)
|
|
}
|
|
if err := cfg.Validate("agent"); err != nil {
|
|
return err
|
|
}
|
|
spool, err := agent.NewSpool(cfg.SpoolRoot, nil)
|
|
if err != nil {
|
|
return err
|
|
}
|
|
report, err := spool.MarkUnknownOnBoot()
|
|
if err != nil {
|
|
return err
|
|
}
|
|
if cfg.GRPCListen == "" {
|
|
return writeResult(map[string]any{
|
|
"role": "agent", "mode": cfg.Mode, "agent_id": cfg.AgentID,
|
|
"version": cfg.Version, "spool": spool.Root(),
|
|
"unknown_executions": report.Unknown, "quarantined": report.Quarantined,
|
|
})
|
|
}
|
|
return serveAgentRPC(cfg, spool, report)
|
|
},
|
|
}
|
|
cmd.Flags().StringVar(&cfg.Mode, "mode", cfg.Mode, "mock, mixed, or real")
|
|
cmd.Flags().StringVar(&cfg.SpoolRoot, "spool", cfg.SpoolRoot, "Agent file spool root")
|
|
cmd.Flags().StringVar(&cfg.AgentID, "agent-id", cfg.AgentID, "stable Agent identifier")
|
|
cmd.Flags().StringVar(&cfg.CellID, "cell-id", cfg.CellID, "stable Cell identifier")
|
|
cmd.Flags().StringVar(&cfg.Version, "version", cfg.Version, "Agent software version")
|
|
cmd.Flags().StringVar(&cfg.StaticArtifactPath, "static-artifact", cfg.StaticArtifactPath, "management-approved static Cell artifact path")
|
|
cmd.Flags().BoolVar(&realCall, "call-once", false, "run one explicit call through the configured transport and AI adapters")
|
|
cmd.Flags().StringVar(&cfg.CallTarget, "call-target", cfg.CallTarget, "raw target number allowed by the static artifact")
|
|
cmd.Flags().StringVar(&cfg.CallTrunkID, "call-trunk-id", cfg.CallTrunkID, "enabled trunk ID from the static artifact")
|
|
cmd.Flags().StringVar(&cfg.CallCallerID, "call-caller-id", cfg.CallCallerID, "deployment-approved caller ID")
|
|
cmd.Flags().StringVar(&cfg.CallAISnapshotPath, "call-ai-snapshot", cfg.CallAISnapshotPath, "immutable AI snapshot path")
|
|
cmd.Flags().IntVar(&cfg.CallMediaPort, "call-media-port", cfg.CallMediaPort, "ExternalMedia UDP port")
|
|
cmd.Flags().StringVar(&cfg.CallRecordingDirectory, "call-recording-dir", cfg.CallRecordingDirectory, "local recording directory")
|
|
return cmd
|
|
}
|
|
|
|
func mediaForCall(artifact contract.StaticCellArtifact, trunkID string) (contract.StaticMedia, error) {
|
|
if artifact.Media == nil {
|
|
return contract.StaticMedia{}, errors.New("static artifact has no default media profile")
|
|
}
|
|
for _, trunk := range artifact.Trunks {
|
|
if trunk.TrunkID != trunkID || !trunk.Enabled {
|
|
continue
|
|
}
|
|
media := *artifact.Media
|
|
if trunk.MediaProfileID == "" {
|
|
return media, nil
|
|
}
|
|
profile, ok := artifact.MediaProfiles[trunk.MediaProfileID]
|
|
if !ok {
|
|
return contract.StaticMedia{}, fmt.Errorf("call trunk %q references unknown media profile %q", trunkID, trunk.MediaProfileID)
|
|
}
|
|
media.Format = profile.Format
|
|
media.SampleRateHz = profile.SampleRateHz
|
|
media.Channels = profile.Channels
|
|
media.PayloadType = profile.PayloadType
|
|
return media, nil
|
|
}
|
|
return contract.StaticMedia{}, fmt.Errorf("call trunk %q is not enabled in the static artifact", trunkID)
|
|
}
|
|
|
|
func buildCallEndpoint(artifact contract.StaticCellArtifact, trunkID, target string) (string, error) {
|
|
if target == "" {
|
|
return "", errors.New("call target is required")
|
|
}
|
|
allowedTarget := false
|
|
for _, allowed := range artifact.AllowedTargets {
|
|
if allowed == target {
|
|
allowedTarget = true
|
|
break
|
|
}
|
|
}
|
|
if !allowedTarget {
|
|
return "", fmt.Errorf("call target is not allowed by the static artifact")
|
|
}
|
|
for _, trunk := range artifact.Trunks {
|
|
if trunk.TrunkID != trunkID || !trunk.Enabled {
|
|
continue
|
|
}
|
|
if trunk.SIPEndpointRef == "" {
|
|
return "", errors.New("selected trunk has no SIP endpoint reference")
|
|
}
|
|
return fmt.Sprintf("PJSIP/%s%s@%s", trunk.DialPrefix, target, trunk.SIPEndpointRef), nil
|
|
}
|
|
return "", fmt.Errorf("call trunk %q is not enabled in the static artifact", trunkID)
|
|
}
|
|
|
|
func validateCallAISnapshot(raw []byte) (ai.Snapshot, error) {
|
|
snapshot, err := ai.Validate(raw)
|
|
if err != nil {
|
|
return ai.Snapshot{}, fmt.Errorf("validate AI snapshot: %w", err)
|
|
}
|
|
return snapshot, nil
|
|
}
|
|
|
|
func runCallOnce(cfg config.Config) error {
|
|
if cfg.StaticArtifactPath == "" || cfg.CallAISnapshotPath == "" {
|
|
return errors.New("--call-once requires --static-artifact and --call-ai-snapshot")
|
|
}
|
|
artifactRaw, err := os.ReadFile(cfg.StaticArtifactPath)
|
|
if err != nil {
|
|
return fmt.Errorf("read static artifact: %w", err)
|
|
}
|
|
artifact, err := contract.ValidateStaticArtifact(artifactRaw, contract.StaticArtifactExpectation{CellID: cfg.CellID, Mode: cfg.Mode})
|
|
if err != nil {
|
|
return fmt.Errorf("validate static artifact: %w", err)
|
|
}
|
|
if cfg.CallTrunkID == "" {
|
|
return errors.New("--call-trunk-id is required")
|
|
}
|
|
trunkBound := false
|
|
for _, trunk := range artifact.Trunks {
|
|
if trunk.TrunkID == cfg.CallTrunkID && trunk.Enabled {
|
|
trunkBound = true
|
|
break
|
|
}
|
|
}
|
|
if !trunkBound {
|
|
return fmt.Errorf("call trunk %q is not enabled in the static artifact", cfg.CallTrunkID)
|
|
}
|
|
if cfg.Mode != "mock" {
|
|
if err := callwindow.Check(time.Now()); err != nil {
|
|
return err
|
|
}
|
|
}
|
|
snapshotRaw, err := os.ReadFile(cfg.CallAISnapshotPath)
|
|
if err != nil {
|
|
return fmt.Errorf("read immutable AI snapshot: %w", err)
|
|
}
|
|
snapshot, err := validateCallAISnapshot(snapshotRaw)
|
|
if err != nil {
|
|
return err
|
|
}
|
|
var snapshotPolicy struct {
|
|
Conversation struct {
|
|
Opening string `json:"opening"`
|
|
MaxTurns int `json:"max_turns"`
|
|
MaxMillis int `json:"max_duration_ms"`
|
|
} `json:"conversation"`
|
|
}
|
|
if err := json.Unmarshal(snapshotRaw, &snapshotPolicy); err != nil {
|
|
return fmt.Errorf("read conversation policy from AI snapshot: %w", err)
|
|
}
|
|
conversation := snapshotPolicy.Conversation
|
|
|
|
var pipeline ai.Pipeline
|
|
switch cfg.Mode {
|
|
case "mock", "mixed":
|
|
pipeline = ai.MockPipeline{MaxAudioBytes: 16 << 20}
|
|
case "real":
|
|
providerCfg, providerErr := ai.LoadProviderPipelineConfigFromEnv()
|
|
if providerErr != nil {
|
|
return providerErr
|
|
}
|
|
pipeline, err = ai.NewProviderPipeline(providerCfg)
|
|
if err != nil {
|
|
return err
|
|
}
|
|
default:
|
|
return fmt.Errorf("unsupported call mode %q", cfg.Mode)
|
|
}
|
|
callCtx, cancel := context.WithTimeout(context.Background(), 4*time.Minute)
|
|
defer cancel()
|
|
if cfg.Mode == "mock" {
|
|
input := make([]byte, 6400)
|
|
flowResult, flowErr := callflow.Execute(callCtx, callflow.NewMemorySession(input), pipeline, snapshot, "您好，这是测试流程。", 10*time.Millisecond)
|
|
if flowErr != nil {
|
|
return flowErr
|
|
}
|
|
return writeResult(map[string]any{
|
|
"role": "agent", "mode": cfg.Mode, "ai_mode": snapshot.Mode, "call": "completed",
|
|
"transcript_chars": len([]rune(flowResult.Turn.Transcript)),
|
|
"reply_chars": len([]rune(flowResult.Turn.Reply)),
|
|
"rtp": map[string]any{"received_packets": flowResult.RTP.ReceivedPackets, "received_bytes": flowResult.RTP.ReceivedBytes, "sent_packets": flowResult.RTP.SentPackets, "sent_bytes": flowResult.RTP.SentBytes},
|
|
})
|
|
}
|
|
if artifact.Media == nil || artifact.ARI == nil || artifact.Recording == nil {
|
|
return errors.New("static artifact is missing transport/media/recording sections")
|
|
}
|
|
if cfg.CallTarget == "" || cfg.ARIUsername == "" || cfg.ARIPassword == "" {
|
|
return errors.New("--call-once requires target and ARI credentials")
|
|
}
|
|
callEndpoint, err := buildCallEndpoint(artifact, cfg.CallTrunkID, cfg.CallTarget)
|
|
if err != nil {
|
|
return err
|
|
}
|
|
selectedMedia, err := mediaForCall(artifact, cfg.CallTrunkID)
|
|
if err != nil {
|
|
return err
|
|
}
|
|
if cfg.CallMediaPort == 0 {
|
|
cfg.CallMediaPort = selectedMedia.Port
|
|
}
|
|
recordingDir := cfg.CallRecordingDirectory
|
|
if recordingDir == "" {
|
|
recordingDir = artifact.Recording.Directory
|
|
}
|
|
maxCallDuration := time.Duration(conversation.MaxMillis) * time.Millisecond
|
|
result, err := callruntime.Run(callCtx, callruntime.Config{
|
|
ARIURL: cfg.ARIURL, ARIWebsocketURL: cfg.ARIWebsocketURL, ARIApplication: artifact.ARI.Application,
|
|
ARIUsername: cfg.ARIUsername, ARIPassword: cfg.ARIPassword, Endpoint: callEndpoint,
|
|
CallerID: cfg.CallCallerID, MediaBind: selectedMedia.BindAddress, MediaPort: cfg.CallMediaPort,
|
|
MediaFormat: selectedMedia.Format, MediaSampleRate: selectedMedia.SampleRateHz,
|
|
PayloadType: uint8(selectedMedia.PayloadType), RecordingDirectory: recordingDir,
|
|
MaxTurns: conversation.MaxTurns, MaxCallDuration: maxCallDuration, OpeningPrompt: conversation.Opening,
|
|
Snapshot: snapshot, Pipeline: pipeline,
|
|
})
|
|
if err != nil {
|
|
return err
|
|
}
|
|
turns := make([]map[string]any, 0, len(result.Turns))
|
|
includeText := os.Getenv("AGENT_CALL_INCLUDE_TEXT") == "1"
|
|
for index, turn := range result.Turns {
|
|
fact := map[string]any{
|
|
"turn": index + 1,
|
|
"transcript_chars": len([]rune(turn.Transcript)),
|
|
"reply_chars": len([]rune(turn.Reply)),
|
|
"invalid_call": turn.InvalidCall,
|
|
"invalid_reason": turn.InvalidReason,
|
|
}
|
|
if includeText {
|
|
fact["transcript"] = turn.Transcript
|
|
fact["reply"] = turn.Reply
|
|
}
|
|
turns = append(turns, fact)
|
|
}
|
|
inboundRecordings := make([]map[string]any, 0, len(result.InboundRecordings))
|
|
for _, recording := range result.InboundRecordings {
|
|
inboundRecordings = append(inboundRecordings, map[string]any{"segment": recording.Segment, "path": recording.Path, "bytes": recording.Bytes, "sha256": recording.SHA256})
|
|
}
|
|
outboundRecordings := make([]map[string]any, 0, len(result.OutboundRecordings))
|
|
for _, recording := range result.OutboundRecordings {
|
|
outboundRecordings = append(outboundRecordings, map[string]any{"segment": recording.Segment, "path": recording.Path, "bytes": recording.Bytes, "sha256": recording.SHA256})
|
|
}
|
|
uploads, err := uploadCallRecordings(callCtx, cfg, result)
|
|
if err != nil {
|
|
return err
|
|
}
|
|
output := map[string]any{
|
|
"role": "agent", "mode": cfg.Mode, "call": "completed", "channel_id": result.ChannelID,
|
|
"transcript_chars": len([]rune(result.Transcript)), "reply_chars": len([]rune(result.Reply)),
|
|
"invalid_call": result.InvalidCall, "invalid_reason": result.InvalidReason,
|
|
"turns": turns,
|
|
"inbound_recording": map[string]any{"path": result.InboundPath, "bytes": result.InboundBytes, "sha256": result.InboundSHA256},
|
|
"outbound_recording": map[string]any{"path": result.OutboundPath, "bytes": result.OutboundBytes, "sha256": result.OutboundSHA256},
|
|
"inbound_recordings": inboundRecordings, "outbound_recordings": outboundRecordings,
|
|
"rtp": map[string]any{"received_packets": result.RTP.ReceivedPackets, "received_bytes": result.RTP.ReceivedBytes, "sent_packets": result.RTP.SentPackets, "sent_bytes": result.RTP.SentBytes},
|
|
}
|
|
if uploads != nil {
|
|
output["oss_uploads"] = uploads
|
|
}
|
|
return writeResult(output)
|
|
}
|
|
|
|
func serveAgentRPC(cfg config.Config, spool *agent.Spool, report agent.RecoveryReport) error {
|
|
var staticArtifactRaw []byte
|
|
if cfg.StaticArtifactPath != "" {
|
|
var err error
|
|
staticArtifactRaw, err = os.ReadFile(cfg.StaticArtifactPath)
|
|
if err != nil {
|
|
return fmt.Errorf("read static Cell artifact: %w", err)
|
|
}
|
|
}
|
|
peerCertificateFingerprints, err := config.ParseCertificateFingerprints(cfg.MTLSPeerCertificateFingerprints)
|
|
if err != nil {
|
|
return fmt.Errorf("parse MTLS_PEER_CERT_FINGERPRINTS: %w", err)
|
|
}
|
|
tlsConfig, err := rpc.LoadServerTLSConfig(cfg.MTLSCAFile, cfg.MTLSCertFile, cfg.MTLSKeyFile)
|
|
if err != nil {
|
|
return err
|
|
}
|
|
var callLogger *calllog.Logger
|
|
if cfg.CallLogPhoneKey != "" {
|
|
path := cfg.CallLogPath
|
|
if path == "" {
|
|
path = filepath.Join(spool.Root(), "call-business.jsonl")
|
|
}
|
|
callLogger, err = calllog.New(path, []byte(cfg.CallLogPhoneKey), nil)
|
|
if err != nil {
|
|
return fmt.Errorf("configure call business log: %w", err)
|
|
}
|
|
}
|
|
listener, err := net.Listen("tcp", cfg.GRPCListen)
|
|
if err != nil {
|
|
return fmt.Errorf("listen Agent gRPC: %w", err)
|
|
}
|
|
defer listener.Close()
|
|
|
|
bootID := fmt.Sprintf("%s-%d", cfg.AgentID, time.Now().UnixNano())
|
|
grpcServer := grpc.NewServer(grpc.Creds(credentials.NewTLS(tlsConfig)))
|
|
handler := rpc.NewServer(rpc.ServerOptions{
|
|
Mode: cfg.Mode,
|
|
Status: &agentv1.AgentStatus{
|
|
AgentId: cfg.AgentID,
|
|
CellId: cfg.CellID,
|
|
BootId: bootID,
|
|
SoftwareVersion: cfg.Version,
|
|
ProtocolVersion: "agent.v1",
|
|
AdmissionState: agentv1.AdmissionState_ADMISSION_STATE_CLOSED,
|
|
StatusReason: fmt.Sprintf("recovered_unknown=%d quarantined=%d spool=%s", len(report.Unknown), len(report.Quarantined), spool.Root()),
|
|
Capabilities: []*agentv1.Capability{{Name: "mode", Value: cfg.Mode}, {Name: "grpc_transport", Value: "unary-mtls"}, {Name: "resource_sample", Value: "partial-unknown"}},
|
|
Resources: health.Sampler{}.Sample(context.Background(), spool.Root()),
|
|
},
|
|
UploadPolicy: &agentv1.UploadPolicy{Enabled: cfg.Mode == "mock", MaxAssetBytes: 16 << 20},
|
|
StaticArtifactRaw: staticArtifactRaw,
|
|
StaticArtifactExpected: contract.StaticArtifactExpectation{CellID: cfg.CellID, Mode: cfg.Mode},
|
|
RequirePeerCertificate: true,
|
|
PeerCertificateFingerprints: peerCertificateFingerprints,
|
|
StatePath: filepath.Join(cfg.SpoolRoot, "rpc-session.json"),
|
|
CallLogger: callLogger,
|
|
})
|
|
agentv1.RegisterAgentControlServiceServer(grpcServer, handler)
|
|
serveCtx, cancel := signalContext()
|
|
defer cancel()
|
|
go func() {
|
|
<-serveCtx.Done()
|
|
grpcServer.GracefulStop()
|
|
}()
|
|
if err := grpcServer.Serve(listener); err != nil && !errors.Is(err, grpc.ErrServerStopped) {
|
|
return err
|
|
}
|
|
return nil
|
|
}
|
|
|
|
type connectedAgentRuntime struct {
|
|
coordinator *dispatcher.AgentCoordinator
|
|
clients []*rpc.Client
|
|
sessions []dispatcher.AgentSession
|
|
}
|
|
|
|
// connectConfiguredAgents performs the Dispatcher startup binding for the
|
|
// deployment-owned endpoint inventory. It is intentionally separate from SaaS
|
|
// commands: tenant input never selects an endpoint or identity.
|
|
func connectConfiguredAgents(ctx context.Context, cfg config.Config) (*connectedAgentRuntime, error) {
|
|
endpoints, err := config.LoadAgentEndpoints(cfg.AgentEndpointsFile)
|
|
if err != nil {
|
|
return nil, err
|
|
}
|
|
runtime := &connectedAgentRuntime{coordinator: dispatcher.NewAgentCoordinator(nil)}
|
|
fail := func(err error) (*connectedAgentRuntime, error) {
|
|
return nil, errors.Join(err, runtime.Close())
|
|
}
|
|
epoch := fmt.Sprintf("%s-%d", cfg.DispatcherID, time.Now().UnixNano())
|
|
for _, endpoint := range endpoints {
|
|
client, dialErr := rpc.DialFromFiles(endpoint.Address, cfg.MTLSCAFile, cfg.MTLSCertFile, cfg.MTLSKeyFile, endpoint.ServerName)
|
|
if dialErr != nil {
|
|
return fail(fmt.Errorf("connect Agent %q: %w", endpoint.AgentID, dialErr))
|
|
}
|
|
runtime.clients = append(runtime.clients, client)
|
|
if registerErr := runtime.coordinator.Register(endpoint.AgentID, client.Agent); registerErr != nil {
|
|
return fail(fmt.Errorf("register Agent %q: %w", endpoint.AgentID, registerErr))
|
|
}
|
|
status, probeErr := runtime.coordinator.Probe(ctx, endpoint.AgentID, endpoint.CellID)
|
|
if probeErr != nil {
|
|
return fail(fmt.Errorf("probe Agent %q: %w", endpoint.AgentID, probeErr))
|
|
}
|
|
// Zero lets Agent-side durable session state allocate the next generation
|
|
// after a Dispatcher restart; hard-coding 1 would self-fence recovery.
|
|
session, activateErr := runtime.coordinator.Activate(ctx, endpoint.AgentID, endpoint.CellID, status.BootId, epoch, 0)
|
|
if activateErr != nil {
|
|
return fail(fmt.Errorf("activate Agent %q: %w", endpoint.AgentID, activateErr))
|
|
}
|
|
runtime.sessions = append(runtime.sessions, session)
|
|
}
|
|
return runtime, nil
|
|
}
|
|
|
|
func (r *connectedAgentRuntime) Close() error {
|
|
if r == nil {
|
|
return nil
|
|
}
|
|
var firstErr error
|
|
for index := len(r.clients) - 1; index >= 0; index-- {
|
|
if err := r.clients[index].Close(); err != nil && firstErr == nil {
|
|
firstErr = err
|
|
}
|
|
}
|
|
return firstErr
|
|
}
|
|
|
|
const dispatcherOutboxFlushInterval = 250 * time.Millisecond
|
|
|
|
func flushDispatcherOutbox(ctx context.Context, d *dispatcher.Dispatcher, batch int, interval time.Duration) {
|
|
flush := func() {
|
|
published, err := d.FlushOutbox(ctx, batch)
|
|
if err != nil {
|
|
if ctx.Err() == nil {
|
|
slog.Error("flush Dispatcher outbox", "error", err)
|
|
}
|
|
return
|
|
}
|
|
if published > 0 {
|
|
slog.Debug("flushed Dispatcher outbox", "published", published)
|
|
}
|
|
}
|
|
flush()
|
|
ticker := time.NewTicker(interval)
|
|
defer ticker.Stop()
|
|
for {
|
|
select {
|
|
case <-ctx.Done():
|
|
return
|
|
case <-ticker.C:
|
|
flush()
|
|
}
|
|
}
|
|
}
|
|
|
|
func newDispatcherCommand() *cobra.Command {
|
|
cfg := config.FromEnv()
|
|
var once bool
|
|
var consume bool
|
|
var tenantKey string
|
|
cmd := &cobra.Command{
|
|
Use: "dispatcher",
|
|
Short: "run the single-active Dispatcher process",
|
|
RunE: func(cmd *cobra.Command, _ []string) error {
|
|
if err := cfg.Validate("dispatcher"); err != nil {
|
|
return err
|
|
}
|
|
st, err := store.Open(cfg.DBPath)
|
|
if err != nil {
|
|
return err
|
|
}
|
|
defer st.Close()
|
|
leaseCtx, cancelLease := signalContext()
|
|
defer cancelLease()
|
|
lease, err := dispatcher.StartLease(leaseCtx, st, "dispatcher-active-"+cfg.DispatcherID, "dispatcher", cfg.DispatcherID, 30*time.Second)
|
|
if err != nil {
|
|
return err
|
|
}
|
|
defer lease.Stop()
|
|
var connectedAgents *connectedAgentRuntime
|
|
if cfg.AgentEndpointsFile != "" {
|
|
startupCtx, cancelStartup := context.WithTimeout(leaseCtx, 15*time.Second)
|
|
connectedAgents, err = connectConfiguredAgents(startupCtx, cfg)
|
|
cancelStartup()
|
|
if err != nil {
|
|
return err
|
|
}
|
|
defer func() {
|
|
if closeErr := connectedAgents.Close(); closeErr != nil {
|
|
slog.Error("close Agent connections", "error", closeErr)
|
|
}
|
|
}()
|
|
}
|
|
var publisher mq.Publisher
|
|
var broker *mq.Broker
|
|
if cfg.RabbitURL != "" {
|
|
broker, err = mq.Open(cfg.RabbitURL, cfg.Exchange)
|
|
if err != nil {
|
|
return err
|
|
}
|
|
defer broker.Close()
|
|
publisher = broker
|
|
}
|
|
d, err := dispatcher.New(st, publisher, nil)
|
|
if err != nil {
|
|
return err
|
|
}
|
|
var dispatcherGRPC *grpc.Server
|
|
var dispatcherListener net.Listener
|
|
if cfg.DispatcherGRPCListen != "" && !once {
|
|
uploadClient, ossErr := ossclient.NewClient(ossclient.Config{
|
|
Endpoint: cfg.OSSEndpoint, Region: cfg.OSSRegion, Bucket: cfg.OSSBucket,
|
|
AccessKeyID: cfg.OSSAccessKeyID, AccessKeySecret: cfg.OSSAccessKeySecret,
|
|
KeyPrefix: cfg.OSSKeyPrefix, GrantTTL: cfg.OSSGrantTTL, MaxAssetBytes: cfg.OSSMaxAssetBytes,
|
|
})
|
|
if ossErr != nil {
|
|
return fmt.Errorf("configure Dispatcher OSS: %w", ossErr)
|
|
}
|
|
peerFingerprints, fingerprintErr := config.ParseCertificateFingerprints(cfg.MTLSPeerCertificateFingerprints)
|
|
if fingerprintErr != nil {
|
|
return fmt.Errorf("parse Dispatcher gRPC mTLS peer fingerprints: %w", fingerprintErr)
|
|
}
|
|
if len(peerFingerprints) == 0 {
|
|
return errors.New("MTLS_PEER_CERT_FINGERPRINTS is required when Dispatcher gRPC is enabled")
|
|
}
|
|
allowedAgentIDs := parseCSVSet(cfg.DispatcherGRPCAllowedAgentIDs)
|
|
uploadHandler, handlerErr := rpc.NewDispatcherUploadServerWithOptions(st, uploadClient, time.Now, rpc.DispatcherUploadOptions{
|
|
RequirePeer: true, PeerCertificateFingerprints: peerFingerprints, AllowedAgentIDs: allowedAgentIDs,
|
|
})
|
|
if handlerErr != nil {
|
|
return handlerErr
|
|
}
|
|
eventHandler, eventErr := rpc.NewDispatcherEventServer(st, rpc.DispatcherEventServerOptions{
|
|
RequirePeer: true, PeerCertificateFingerprints: peerFingerprints,
|
|
AllowedAgentIDs: allowedAgentIDs, Now: time.Now,
|
|
})
|
|
if eventErr != nil {
|
|
return eventErr
|
|
}
|
|
dispatcherHandler := rpc.NewDispatcherServer(uploadHandler, eventHandler)
|
|
tlsConfig, tlsErr := rpc.LoadServerTLSConfig(cfg.MTLSCAFile, cfg.MTLSCertFile, cfg.MTLSKeyFile)
|
|
if tlsErr != nil {
|
|
return fmt.Errorf("configure Dispatcher gRPC mTLS: %w", tlsErr)
|
|
}
|
|
dispatcherListener, err = net.Listen("tcp", cfg.DispatcherGRPCListen)
|
|
if err != nil {
|
|
return fmt.Errorf("listen Dispatcher gRPC: %w", err)
|
|
}
|
|
dispatcherGRPC = grpc.NewServer(grpc.Creds(credentials.NewTLS(tlsConfig)))
|
|
agentv1.RegisterAgentControlServiceServer(dispatcherGRPC, dispatcherHandler)
|
|
go func() {
|
|
if serveErr := dispatcherGRPC.Serve(dispatcherListener); serveErr != nil && !errors.Is(serveErr, grpc.ErrServerStopped) {
|
|
slog.Error("Dispatcher gRPC stopped", "error", serveErr)
|
|
}
|
|
}()
|
|
defer func() {
|
|
dispatcherGRPC.GracefulStop()
|
|
_ = dispatcherListener.Close()
|
|
}()
|
|
}
|
|
result := map[string]any{"role": "dispatcher", "mode": cfg.Mode, "db": cfg.DBPath}
|
|
if dispatcherGRPC != nil {
|
|
result["dispatcher_grpc"] = cfg.DispatcherGRPCListen
|
|
}
|
|
if connectedAgents != nil {
|
|
sessions := make([]map[string]any, 0, len(connectedAgents.sessions))
|
|
for _, session := range connectedAgents.sessions {
|
|
sessions = append(sessions, map[string]any{
|
|
"agent_id": session.AgentID, "cell_id": session.CellID,
|
|
"boot_id": session.BootID, "session_generation": session.SessionGeneration,
|
|
})
|
|
}
|
|
result["agent_sessions"] = sessions
|
|
}
|
|
if once && consume {
|
|
return errors.New("--once cannot be combined with --consume")
|
|
}
|
|
if consume {
|
|
if broker == nil || tenantKey == "" {
|
|
return errors.New("--consume requires --rabbit-url and --tenant-key")
|
|
}
|
|
if cfg.ControlListen != "" {
|
|
return errors.New("--consume cannot be combined with --control-listen")
|
|
}
|
|
flushCtx, cancelFlush := context.WithCancel(leaseCtx)
|
|
flushDone := make(chan struct{})
|
|
go func() {
|
|
defer close(flushDone)
|
|
flushDispatcherOutbox(flushCtx, d, cfg.OutboxBatch, dispatcherOutboxFlushInterval)
|
|
}()
|
|
defer func() {
|
|
cancelFlush()
|
|
<-flushDone
|
|
}()
|
|
if err := d.ConsumeTenant(leaseCtx, broker, tenantKey); err != nil && !errors.Is(err, context.Canceled) {
|
|
return err
|
|
}
|
|
return nil
|
|
}
|
|
if once {
|
|
if publisher == nil {
|
|
return errors.New("--once requires RABBITMQ_URL or a configured publisher")
|
|
}
|
|
count, err := d.FlushOutbox(cmd.Context(), cfg.OutboxBatch)
|
|
if err != nil {
|
|
return err
|
|
}
|
|
result["published"] = count
|
|
}
|
|
if cfg.ControlListen == "" {
|
|
if dispatcherGRPC == nil {
|
|
return writeResult(result)
|
|
}
|
|
select {
|
|
case <-leaseCtx.Done():
|
|
return nil
|
|
case err := <-lease.Lost():
|
|
return fmt.Errorf("dispatcher lease lost: %w", err)
|
|
}
|
|
}
|
|
if once {
|
|
return errors.New("--once cannot be combined with --control-listen")
|
|
}
|
|
server := &http.Server{Addr: cfg.ControlListen, Handler: control.Handler{Store: st, BearerToken: cfg.ControlToken}}
|
|
go func() {
|
|
select {
|
|
case <-leaseCtx.Done():
|
|
case <-lease.Lost():
|
|
}
|
|
shutdownCtx, shutdownCancel := context.WithTimeout(context.Background(), 5*time.Second)
|
|
defer shutdownCancel()
|
|
_ = server.Shutdown(shutdownCtx)
|
|
}()
|
|
if err := server.ListenAndServe(); !errors.Is(err, http.ErrServerClosed) {
|
|
return err
|
|
}
|
|
select {
|
|
case err := <-lease.Lost():
|
|
return fmt.Errorf("dispatcher lease lost: %w", err)
|
|
default:
|
|
return nil
|
|
}
|
|
},
|
|
}
|
|
cmd.Flags().StringVar(&cfg.Mode, "mode", cfg.Mode, "mock, mixed, or real")
|
|
cmd.Flags().StringVar(&cfg.DBPath, "db", cfg.DBPath, "Dispatcher SQLite path")
|
|
cmd.Flags().StringVar(&cfg.DispatcherID, "dispatcher-id", cfg.DispatcherID, "single-active Dispatcher holder identity")
|
|
cmd.Flags().StringVar(&cfg.RabbitURL, "rabbit-url", cfg.RabbitURL, "RabbitMQ URL")
|
|
cmd.Flags().StringVar(&cfg.Exchange, "exchange", cfg.Exchange, "durable command exchange")
|
|
cmd.Flags().IntVar(&cfg.OutboxBatch, "outbox-batch", cfg.OutboxBatch, "maximum outbox messages per run")
|
|
cmd.Flags().StringVar(&cfg.ControlListen, "control-listen", cfg.ControlListen, "internal control HTTP listen address; empty disables server")
|
|
cmd.Flags().StringVar(&cfg.ControlToken, "control-token", cfg.ControlToken, "bearer token for internal control HTTP")
|
|
cmd.Flags().StringVar(&cfg.AgentEndpointsFile, "agent-endpoints-file", cfg.AgentEndpointsFile, "strict JSON file of Dispatcher-owned Agent endpoints")
|
|
cmd.Flags().StringVar(&tenantKey, "tenant-key", "", "tenant key to consume from its command queue")
|
|
cmd.Flags().BoolVar(&consume, "consume", false, "consume one tenant command queue")
|
|
cmd.Flags().BoolVar(&once, "once", false, "flush one outbox batch and exit")
|
|
return cmd
|
|
}
|
|
|
|
func parseCSVSet(raw string) map[string]struct{} {
|
|
result := make(map[string]struct{})
|
|
for _, item := range strings.Split(raw, ",") {
|
|
item = strings.TrimSpace(item)
|
|
if item != "" {
|
|
result[item] = struct{}{}
|
|
}
|
|
}
|
|
return result
|
|
}
|
|
|
|
func marshalResultForTest(value any) ([]byte, error) {
|
|
return json.Marshal(value)
|
|
}
|
|
|
|
func writeResult(value any) error {
|
|
data, err := marshalResultForTest(value)
|
|
if err != nil {
|
|
return err
|
|
}
|
|
_, err = os.Stdout.Write(append(data, '\n'))
|
|
return err
|
|
}
|
|
|
|
func signalContext() (context.Context, context.CancelFunc) {
|
|
return signal.NotifyContext(context.Background(), syscall.SIGINT, syscall.SIGTERM)
|
|
}
|