Merge pull request #890 from chaitin/fix/mcp-sync-health-status

同步 MCP 工具时更新健康状态

Author: yokowu

backend/biz/mcphub/control/syncer/service.go

@@ -16,6 +16,7 @@ type upstreamRepo interface {
 	Get(ctx context.Context, id uuid.UUID) (*repo.UpstreamConfig, error)
 	MarkSyncSuccess(ctx context.Context, id uuid.UUID) error
 	MarkSyncFailed(ctx context.Context, id uuid.UUID) error
+	MarkHealthStatus(ctx context.Context, id uuid.UUID, healthy bool) error
 }
 
 type toolRepo interface {
@@ -54,6 +55,11 @@ func (s *Service) Sync(ctx context.Context, upstreamID uuid.UUID) error {
 
 	tools, err := s.client.ListTools(ctx, upstream)
 	if err != nil {
+		_ = s.upstreams.MarkHealthStatus(ctx, upstreamID, false)
+		_ = s.upstreams.MarkSyncFailed(ctx, upstreamID)
+		return err
+	}
+	if err := s.upstreams.MarkHealthStatus(ctx, upstreamID, true); err != nil {
 		_ = s.upstreams.MarkSyncFailed(ctx, upstreamID)
 		return err
 	}

backend/biz/mcphub/control/syncer/service_test.go

@@ -0,0 +1,98 @@
+package syncer
+
+import (
+	"context"
+	"errors"
+	"testing"
+
+	"github.com/google/uuid"
+
+	"github.com/chaitin/MonkeyCode/backend/biz/mcphub/repo"
+)
+
+func TestSyncMarksUpstreamHealthy(t *testing.T) {
+	upstreamID := uuid.New()
+	upstreams := &upstreamRepoStub{
+		upstream: &repo.UpstreamConfig{ID: upstreamID, Slug: "images"},
+	}
+	service := NewService(upstreams, &toolRepoStub{}, &registryStub{}, &upstreamClientStub{
+		tools: []repo.UpstreamTool{{Name: "search"}},
+	})
+
+	if err := service.Sync(context.Background(), upstreamID); err != nil {
+		t.Fatal(err)
+	}
+	if upstreams.healthy == nil || !*upstreams.healthy {
+		t.Fatalf("health status = %v, want healthy", upstreams.healthy)
+	}
+	if !upstreams.syncSuccess {
+		t.Fatal("sync success was not recorded")
+	}
+}
+
+func TestSyncMarksUpstreamUnhealthyWhenListToolsFails(t *testing.T) {
+	upstreamID := uuid.New()
+	upstreams := &upstreamRepoStub{
+		upstream: &repo.UpstreamConfig{ID: upstreamID, Slug: "images"},
+	}
+	service := NewService(upstreams, &toolRepoStub{}, &registryStub{}, &upstreamClientStub{
+		err: errors.New("upstream unavailable"),
+	})
+
+	if err := service.Sync(context.Background(), upstreamID); err == nil {
+		t.Fatal("Sync() error = nil, want upstream error")
+	}
+	if upstreams.healthy == nil || *upstreams.healthy {
+		t.Fatalf("health status = %v, want unhealthy", upstreams.healthy)
+	}
+	if !upstreams.syncFailed {
+		t.Fatal("sync failure was not recorded")
+	}
+}
+
+type upstreamRepoStub struct {
+	upstream    *repo.UpstreamConfig
+	healthy     *bool
+	syncSuccess bool
+	syncFailed  bool
+}
+
+func (s *upstreamRepoStub) Get(context.Context, uuid.UUID) (*repo.UpstreamConfig, error) {
+	return s.upstream, nil
+}
+
+func (s *upstreamRepoStub) MarkSyncSuccess(context.Context, uuid.UUID) error {
+	s.syncSuccess = true
+	return nil
+}
+
+func (s *upstreamRepoStub) MarkSyncFailed(context.Context, uuid.UUID) error {
+	s.syncFailed = true
+	return nil
+}
+
+func (s *upstreamRepoStub) MarkHealthStatus(_ context.Context, _ uuid.UUID, healthy bool) error {
+	s.healthy = &healthy
+	return nil
+}
+
+type toolRepoStub struct{}
+
+func (s *toolRepoStub) ReplaceByUpstream(context.Context, uuid.UUID, []repo.UpsertToolInput) error {
+	return nil
+}
+
+type registryStub struct{}
+
+func (s *registryStub) Publish(context.Context) error {
+	return nil
+}
+
+type upstreamClientStub struct {
+	tools []repo.UpstreamTool
+	err   error
+}
+
+func (s *upstreamClientStub) ListTools(context.Context, *repo.UpstreamConfig) ([]repo.UpstreamTool, error) {
+	return s.tools, s.err
+}

backend/biz/mcphub/repo/upstream.go

@@ -93,6 +93,20 @@ func (r *UpstreamRepo) MarkSyncFailed(ctx context.Context, id uuid.UUID) error {
 	return nil
 }
 
+func (r *UpstreamRepo) MarkHealthStatus(ctx context.Context, id uuid.UUID, healthy bool) error {
+	status := "unhealthy"
+	if healthy {
+		status = "healthy"
+	}
+	if err := r.db.MCPUpstream.UpdateOneID(id).
+		SetHealthStatus(status).
+		SetHealthCheckedAt(time.Now()).
+		Exec(ctx); err != nil {
+		return fmt.Errorf("mark upstream %s health status: %w", id, err)
+	}
+	return nil
+}
+
 func IsUpstreamNotFound(err error) bool {
 	return errors.Is(err, ErrUpstreamNotFound)
 }

backend/biz/mcphub/repo/upstream_test.go

@@ -0,0 +1,44 @@
+package repo
+
+import (
+	"context"
+	"testing"
+
+	"github.com/google/uuid"
+	_ "github.com/mattn/go-sqlite3"
+
+	"github.com/chaitin/MonkeyCode/backend/db/enttest"
+)
+
+func TestMarkHealthStatus(t *testing.T) {
+	client := enttest.Open(t, "sqlite3", "file:mcphub-upstream-health?mode=memory&cache=shared&_fk=1")
+	defer client.Close()
+	ctx := context.Background()
+	upstreamID := uuid.New()
+
+	client.MCPUpstream.Create().
+		SetID(upstreamID).
+		SetName("images").
+		SetSlug("images").
+		SetScope("platform").
+		SetType("server").
+		SetURL("https://example.com/mcp").
+		SaveX(ctx)
+
+	repo := NewUpstreamRepo(client)
+	if err := repo.MarkHealthStatus(ctx, upstreamID, true); err != nil {
+		t.Fatal(err)
+	}
+	row := client.MCPUpstream.GetX(ctx, upstreamID)
+	if row.HealthStatus != "healthy" || row.HealthCheckedAt == nil {
+		t.Fatalf("healthy upstream = %+v", row)
+	}
+
+	if err := repo.MarkHealthStatus(ctx, upstreamID, false); err != nil {
+		t.Fatal(err)
+	}
+	row = client.MCPUpstream.GetX(ctx, upstreamID)
+	if row.HealthStatus != "unhealthy" || row.HealthCheckedAt == nil {
+		t.Fatalf("unhealthy upstream = %+v", row)
+	}
+}