[feat][backend] volcengine agentkit

backend/cmd/main.go

@@ -9,6 +9,8 @@ import (
 	"fmt"
 	"net/url"
 	"os"
+	"os/signal"
+	"syscall"
 	"time"
 
 	"github.com/bytedance/gg/gptr"
@@ -60,12 +62,21 @@ func main() {
 	if err := initTracer(handler); err != nil {
 		panic(err)
 	}
-	consumerWorkers := MustInitConsumerWorkers(c.cfgFactory, handler, handler, handler, handler)
-	if err := registry.NewConsumerRegistry(c.mqFactory).Register(consumerWorkers).StartAll(ctx); err != nil {
+
+	signalCtx, signalCancel := signal.NotifyContext(ctx, syscall.SIGTERM, syscall.SIGINT)
+	defer signalCancel()
+
+	r := registry.NewConsumerRegistryWithShutdown(signalCtx, c.mqFactory).Register(MustInitConsumerWorkers(c.cfgFactory, handler, handler, handler, handler))
+	if err := r.StartAll(ctx); err != nil {
 		panic(err)
 	}
 
-	api.Start(handler)
+	go api.Start(handler)
+	<-signalCtx.Done()
+
+	stopCtx, stopCancel := context.WithTimeout(context.Background(), 30*time.Second)
+	defer stopCancel()
+	_ = r.StopAll(stopCtx)
 }
 
 type ComponentConfig struct {

backend/infra/mq/factory.go

@@ -31,6 +31,9 @@ type ProducerConfig struct {
 	FlushFrequency time.Duration
 	// How long to wait for the cluster to settle between retries
 	RetryBackoff time.Duration
+
+	AccessKey    *string
+	AccessSecret *string
 }
 
 type ConsumerConfig struct {
@@ -50,6 +53,9 @@ type ConsumerConfig struct {
 	ConsumeTimeout time.Duration
 	EnablePPE      *bool
 	IsEnabled      *bool
+
+	AccessKey    *string
+	AccessSecret *string
 }
 
 type CompressionCodec int

backend/infra/mq/registry.go

@@ -12,6 +12,7 @@ import (
 type ConsumerRegistry interface {
 	Register(worker []IConsumerWorker) ConsumerRegistry
 	StartAll(ctx context.Context) error
+	StopAll(ctx context.Context) error
 }
 
 type IConsumerWorker interface {

backend/infra/mq/registry/registry.go

@@ -5,6 +5,7 @@ package registry
 
 import (
 	"context"
+	"errors"
 
 	"github.com/coze-dev/coze-loop/backend/infra/mq"
 	"github.com/coze-dev/coze-loop/backend/pkg/errorx"
@@ -14,20 +15,27 @@ import (
 )
 
 type defaultConsumerRegistry struct {
-	factory mq.IFactory
-	workers []mq.IConsumerWorker
+	factory     mq.IFactory
+	workers     []mq.IConsumerWorker
+	consumers   []mq.IConsumer
+	shutdownCtx context.Context
 }
 
 func NewConsumerRegistry(factory mq.IFactory) mq.ConsumerRegistry {
 	return &defaultConsumerRegistry{factory: factory}
 }
 
+func NewConsumerRegistryWithShutdown(shutdownCtx context.Context, factory mq.IFactory) mq.ConsumerRegistry {
+	return &defaultConsumerRegistry{factory: factory, shutdownCtx: shutdownCtx}
+}
+
 func (d *defaultConsumerRegistry) Register(worker []mq.IConsumerWorker) mq.ConsumerRegistry {
 	d.workers = append(d.workers, worker...)
 	return d
 }
 
 func (d *defaultConsumerRegistry) StartAll(ctx context.Context) error {
+	d.consumers = nil
 	for _, worker := range d.workers {
 		cfg, err := worker.ConsumerCfg(ctx)
 		if err != nil {
@@ -39,14 +47,47 @@ func (d *defaultConsumerRegistry) StartAll(ctx context.Context) error {
 			return errorx.Wrapf(err, "NewConsumer fail, cfg: %v", json.Jsonify(cfg))
 		}
 
-		consumer.RegisterHandler(newSafeConsumerWrapper(worker))
+		handler := newSafeConsumerWrapper(worker)
+		if d.shutdownCtx != nil {
+			handler = newShutdownContextWrapper(handler, d.shutdownCtx)
+		}
+		consumer.RegisterHandler(handler)
 		if err := consumer.Start(); err != nil {
 			return errorx.Wrapf(err, "StartConsumer fail, cfg: %v", json.Jsonify(cfg))
 		}
+		d.consumers = append(d.consumers, consumer)
 	}
 	return nil
 }
 
+func (d *defaultConsumerRegistry) StopAll(ctx context.Context) error {
+	if len(d.consumers) == 0 {
+		return nil
+	}
+	var errs []error
+	for i := len(d.consumers) - 1; i >= 0; i-- {
+		select {
+		case <-ctx.Done():
+			errs = append(errs, ctx.Err())
+			return errors.Join(errs...)
+		default:
+			consumer := d.consumers[i]
+			done := make(chan error, 1)
+			go func(c mq.IConsumer) { done <- c.Close() }(consumer)
+			select {
+			case err := <-done:
+				if err != nil {
+					errs = append(errs, err)
+				}
+			case <-ctx.Done():
+				errs = append(errs, ctx.Err())
+				return errors.Join(errs...)
+			}
+		}
+	}
+	return errors.Join(errs...)
+}
+
 type safeConsumerHandlerDecorator struct {
 	handler mq.IConsumerHandler
 }
@@ -59,3 +100,25 @@ func (s *safeConsumerHandlerDecorator) HandleMessage(ctx context.Context, msg *m
 func newSafeConsumerWrapper(h mq.IConsumerHandler) mq.IConsumerHandler {
 	return &safeConsumerHandlerDecorator{handler: h}
 }
+
+type shutdownContextDecorator struct {
+	handler     mq.IConsumerHandler
+	shutdownCtx context.Context
+}
+
+func (s *shutdownContextDecorator) HandleMessage(ctx context.Context, msg *mq.MessageExt) error {
+	nctx, cancel := context.WithCancel(ctx)
+	go func() {
+		defer goroutine.Recovery(ctx)
+		select {
+		case <-ctx.Done():
+		case <-s.shutdownCtx.Done():
+		}
+		cancel()
+	}()
+	return s.handler.HandleMessage(nctx, msg)
+}
+
+func newShutdownContextWrapper(h mq.IConsumerHandler, shutdownCtx context.Context) mq.IConsumerHandler {
+	return &shutdownContextDecorator{handler: h, shutdownCtx: shutdownCtx}
+}

backend/infra/mq/registry/registry_test.go

@@ -20,6 +20,7 @@ func TestDefaultConsumerRegistry_StartAll(t *testing.T) {
 		name          string
 		workers       []mq.IConsumerWorker
 		setupMocks    func(*mocks.MockIFactory, []*mocks.MockIConsumer, []*mocks.MockIConsumerWorker)
+		shutdownCtx   context.Context
 		expectedError error
 	}{
 		{
@@ -39,6 +40,21 @@ func TestDefaultConsumerRegistry_StartAll(t *testing.T) {
 			},
 			expectedError: nil,
 		},
+		{
+			name: "successfully start all workers with shutdown ctx",
+			workers: []mq.IConsumerWorker{
+				mocks.NewMockIConsumerWorker(gomock.NewController(t)),
+			},
+			setupMocks: func(factory *mocks.MockIFactory, consumers []*mocks.MockIConsumer, workers []*mocks.MockIConsumerWorker) {
+				cfg := &mq.ConsumerConfig{}
+				workers[0].EXPECT().ConsumerCfg(gomock.Any()).Return(cfg, nil)
+				consumers[0].EXPECT().RegisterHandler(gomock.Any()).Return()
+				consumers[0].EXPECT().Start().Return(nil)
+				factory.EXPECT().NewConsumer(gomock.Any()).Return(consumers[0], nil)
+			},
+			shutdownCtx:   context.Background(),
+			expectedError: nil,
+		},
 		{
 			name: "fail to get consumer config",
 			workers: []mq.IConsumerWorker{
@@ -92,9 +108,12 @@ func TestDefaultConsumerRegistry_StartAll(t *testing.T) {
 			}
 
 			tt.setupMocks(factory, consumers, workers)
-
-			registry := NewConsumerRegistry(factory).Register(tt.workers)
-
+			var registry mq.ConsumerRegistry
+			if tt.shutdownCtx != nil {
+				registry = NewConsumerRegistryWithShutdown(tt.shutdownCtx, factory).Register(tt.workers)
+			} else {
+				registry = NewConsumerRegistry(factory).Register(tt.workers)
+			}
 			err := registry.StartAll(context.Background())
 			if tt.expectedError != nil {
 				assert.Error(t, err)
@@ -106,6 +125,69 @@ func TestDefaultConsumerRegistry_StartAll(t *testing.T) {
 	}
 }
 
+func TestDefaultConsumerRegistry_StopAll(t *testing.T) {
+	t.Run("no consumers", func(t *testing.T) {
+		ctrl := gomock.NewController(t)
+		defer ctrl.Finish()
+		factory := mocks.NewMockIFactory(ctrl)
+		registry := NewConsumerRegistry(factory)
+		err := registry.StopAll(context.Background())
+		assert.NoError(t, err)
+	})
+
+	t.Run("successfully stop all consumers", func(t *testing.T) {
+		ctrl := gomock.NewController(t)
+		defer ctrl.Finish()
+		factory := mocks.NewMockIFactory(ctrl)
+		workers := []mq.IConsumerWorker{
+			mocks.NewMockIConsumerWorker(ctrl),
+			mocks.NewMockIConsumerWorker(ctrl),
+		}
+		consumers := []*mocks.MockIConsumer{
+			mocks.NewMockIConsumer(ctrl),
+			mocks.NewMockIConsumer(ctrl),
+		}
+		cfg := &mq.ConsumerConfig{}
+		for i := range workers {
+			workers[i].(*mocks.MockIConsumerWorker).EXPECT().ConsumerCfg(gomock.Any()).Return(cfg, nil)
+			factory.EXPECT().NewConsumer(gomock.Any()).Return(consumers[i], nil)
+			consumers[i].EXPECT().RegisterHandler(gomock.Any())
+			consumers[i].EXPECT().Start().Return(nil)
+		}
+		registry := NewConsumerRegistry(factory).Register(workers)
+		err := registry.StartAll(context.Background())
+		assert.NoError(t, err)
+
+		// StopAll 按逆序关闭，先关 consumers[1] 再关 consumers[0]
+		consumers[1].EXPECT().Close().Return(nil)
+		consumers[0].EXPECT().Close().Return(nil)
+		err = registry.StopAll(context.Background())
+		assert.NoError(t, err)
+	})
+
+	t.Run("context cancelled during stop", func(t *testing.T) {
+		ctrl := gomock.NewController(t)
+		defer ctrl.Finish()
+		factory := mocks.NewMockIFactory(ctrl)
+		worker := mocks.NewMockIConsumerWorker(ctrl)
+		consumer := mocks.NewMockIConsumer(ctrl)
+		cfg := &mq.ConsumerConfig{}
+		worker.EXPECT().ConsumerCfg(gomock.Any()).Return(cfg, nil)
+		factory.EXPECT().NewConsumer(gomock.Any()).Return(consumer, nil)
+		consumer.EXPECT().RegisterHandler(gomock.Any())
+		consumer.EXPECT().Start().Return(nil)
+		registry := NewConsumerRegistry(factory).Register([]mq.IConsumerWorker{worker})
+		err := registry.StartAll(context.Background())
+		assert.NoError(t, err)
+
+		ctx, cancel := context.WithCancel(context.Background())
+		cancel()
+		err = registry.StopAll(ctx)
+		assert.Error(t, err)
+		assert.True(t, errors.Is(err, context.Canceled))
+	})
+}
+
 func TestSafeConsumerHandlerDecorator_HandleMessage(t *testing.T) {
 	tests := []struct {
 		name          string
@@ -150,3 +232,72 @@ func TestSafeConsumerHandlerDecorator_HandleMessage(t *testing.T) {
 		})
 	}
 }
+
+func TestNewConsumerRegistryWithShutdown(t *testing.T) {
+	ctrl := gomock.NewController(t)
+	defer ctrl.Finish()
+
+	factory := mocks.NewMockIFactory(ctrl)
+	shutdownCtx, cancel := context.WithCancel(context.Background())
+	defer cancel()
+
+	registry := NewConsumerRegistryWithShutdown(shutdownCtx, factory).(*defaultConsumerRegistry)
+	assert.Equal(t, factory, registry.factory)
+	assert.Equal(t, shutdownCtx, registry.shutdownCtx)
+}
+
+func TestShutdownContextDecorator_HandleMessage(t *testing.T) {
+	ctrl := gomock.NewController(t)
+	defer ctrl.Finish()
+
+	mockHandler := mocks.NewMockIConsumerWorker(ctrl)
+	shutdownCtx, shutdownCancel := context.WithCancel(context.Background())
+
+	decorator := &shutdownContextDecorator{
+		handler:     mockHandler,
+		shutdownCtx: shutdownCtx,
+	}
+
+	tests := []struct {
+		name          string
+		setupMock     func()
+		triggerCancel func()
+		ctx           context.Context
+	}{
+		{
+			name: "normal execution",
+			setupMock: func() {
+				mockHandler.EXPECT().HandleMessage(gomock.Any(), gomock.Any()).Return(nil)
+			},
+			triggerCancel: func() {},
+			ctx:           context.Background(),
+		},
+		{
+			name: "shutdown context cancelled",
+			setupMock: func() {
+				mockHandler.EXPECT().HandleMessage(gomock.Any(), gomock.Any()).DoAndReturn(func(ctx context.Context, msg *mq.MessageExt) error {
+					<-ctx.Done()
+					return ctx.Err()
+				})
+			},
+			triggerCancel: func() {
+				shutdownCancel()
+			},
+			ctx: context.Background(),
+		},
+	}
+
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			tt.setupMock()
+			go tt.triggerCancel()
+			err := decorator.HandleMessage(tt.ctx, &mq.MessageExt{})
+			if tt.name == "shutdown context cancelled" {
+				assert.Error(t, err)
+				assert.True(t, errors.Is(err, context.Canceled))
+			} else {
+				assert.NoError(t, err)
+			}
+		})
+	}
+}