350 lines
10 KiB
Go
350 lines
10 KiB
Go
// Copyright 2026 Alibaba Group
|
|
// Licensed under the Apache License, Version 2.0 (the "License");
|
|
// you may not use this file except in compliance with the License.
|
|
// You may obtain a copy of the License at
|
|
//
|
|
// http://www.apache.org/licenses/LICENSE-2.0
|
|
//
|
|
// Unless required by applicable law or agreed to in writing, software
|
|
// distributed under the License is distributed on an "AS IS" BASIS,
|
|
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
|
// See the License for the specific language governing permissions and
|
|
// limitations under the License.
|
|
|
|
package discovery
|
|
|
|
import (
|
|
"context"
|
|
"encoding/json"
|
|
"errors"
|
|
"fmt"
|
|
"log/slog"
|
|
"os"
|
|
"strings"
|
|
|
|
"github.com/DingTalk-Real-AI/dingtalk-workspace-cli/internal/cache"
|
|
"github.com/DingTalk-Real-AI/dingtalk-workspace-cli/internal/market"
|
|
"github.com/DingTalk-Real-AI/dingtalk-workspace-cli/internal/transport"
|
|
)
|
|
|
|
const (
|
|
tenantEnv = "DWS_TENANT"
|
|
authIdentityEnv = "DWS_AUTH_IDENTITY"
|
|
)
|
|
|
|
var errCLIServerSkipped = errors.New("server marked cli.skip")
|
|
|
|
type Service struct {
|
|
MarketClient *market.Client
|
|
Transport *transport.Client
|
|
Cache *cache.Store
|
|
Tenant string
|
|
AuthIdentity string
|
|
Logger *slog.Logger
|
|
}
|
|
|
|
type RuntimeServer struct {
|
|
Server market.ServerDescriptor `json:"server"`
|
|
NegotiatedProtocolVersion string `json:"negotiated_protocol_version"`
|
|
Tools []transport.ToolDescriptor `json:"tools"`
|
|
Source string `json:"source"`
|
|
Degraded bool `json:"degraded"`
|
|
}
|
|
|
|
type RuntimeFailure struct {
|
|
ServerKey string
|
|
Err error
|
|
}
|
|
|
|
func NewService(marketClient *market.Client, transportClient *transport.Client, cacheStore *cache.Store) *Service {
|
|
return &Service{
|
|
MarketClient: marketClient,
|
|
Transport: transportClient,
|
|
Cache: cacheStore,
|
|
Tenant: resolveTenant(),
|
|
AuthIdentity: resolveAuthIdentity(),
|
|
}
|
|
}
|
|
|
|
func (s *Service) DiscoverServers(ctx context.Context) ([]market.ServerDescriptor, error) {
|
|
partition := s.partition()
|
|
|
|
response, err := s.MarketClient.FetchServers(ctx, 200)
|
|
if err == nil {
|
|
servers := market.NormalizeServers(response, "live_market")
|
|
_ = s.Cache.SaveRegistry(partition, cache.RegistrySnapshot{Servers: servers})
|
|
return servers, nil
|
|
}
|
|
|
|
snapshot, freshness, cacheErr := s.Cache.LoadRegistry(partition)
|
|
if cacheErr == nil {
|
|
servers := append([]market.ServerDescriptor(nil), snapshot.Servers...)
|
|
for idx := range servers {
|
|
servers[idx].Source = string(freshness) + "_cache"
|
|
servers[idx].Degraded = true
|
|
}
|
|
return servers, nil
|
|
}
|
|
|
|
return nil, fmt.Errorf("discover servers: market fetch failed and no cache available: %w", err)
|
|
}
|
|
|
|
func (s *Service) DiscoverServerRuntime(ctx context.Context, server market.ServerDescriptor) (RuntimeServer, error) {
|
|
if server.CLI.Skip {
|
|
return RuntimeServer{}, errCLIServerSkipped
|
|
}
|
|
|
|
partition := s.partition()
|
|
|
|
initialize, err := s.Transport.Initialize(ctx, server.Endpoint)
|
|
if err == nil {
|
|
// NotifyInitialized is best-effort; log but do not fail on error.
|
|
if notifyErr := s.Transport.NotifyInitialized(ctx, server.Endpoint); notifyErr != nil && s.Logger != nil {
|
|
s.Logger.Debug("NotifyInitialized failed", "server", server.Key, "error", notifyErr)
|
|
}
|
|
tools, listErr := s.Transport.ListTools(ctx, server.Endpoint)
|
|
if listErr == nil {
|
|
runtimeTools := append([]transport.ToolDescriptor(nil), tools.Tools...)
|
|
var actionVersions map[string]string
|
|
if server.DetailLocator.MCPID > 0 {
|
|
if detail, detailErr := s.DiscoverDetail(ctx, server); detailErr == nil {
|
|
runtimeTools = mergeRuntimeToolsWithDetail(runtimeTools, detail)
|
|
actionVersions = cache.ExtractActionVersions(detail.Result.Tools)
|
|
}
|
|
}
|
|
_ = s.Cache.SaveTools(partition, server.Key, cache.ToolsSnapshot{
|
|
ServerKey: server.Key,
|
|
ProtocolVersion: initialize.ProtocolVersion,
|
|
Tools: runtimeTools,
|
|
ActionVersions: actionVersions,
|
|
})
|
|
server.NegotiatedProtocolVersion = initialize.ProtocolVersion
|
|
server.Source = "live_runtime"
|
|
return RuntimeServer{
|
|
Server: server,
|
|
NegotiatedProtocolVersion: initialize.ProtocolVersion,
|
|
Tools: runtimeTools,
|
|
Source: "live_runtime",
|
|
Degraded: false,
|
|
}, nil
|
|
}
|
|
err = listErr
|
|
}
|
|
|
|
if errors.Is(err, context.Canceled) || errors.Is(err, context.DeadlineExceeded) {
|
|
return RuntimeServer{}, err
|
|
}
|
|
|
|
snapshot, freshness, cacheErr := s.Cache.LoadTools(partition, server.Key)
|
|
if cacheErr != nil {
|
|
return RuntimeServer{}, fmt.Errorf("server %s: runtime discovery failed and no cache available: %w", server.Key, err)
|
|
}
|
|
|
|
server.NegotiatedProtocolVersion = snapshot.ProtocolVersion
|
|
server.Source = string(freshness) + "_cache"
|
|
server.Degraded = true
|
|
return RuntimeServer{
|
|
Server: server,
|
|
NegotiatedProtocolVersion: snapshot.ProtocolVersion,
|
|
Tools: snapshot.Tools,
|
|
Source: string(freshness) + "_cache",
|
|
Degraded: true,
|
|
}, nil
|
|
}
|
|
|
|
func (s *Service) DiscoverAllRuntime(ctx context.Context, servers []market.ServerDescriptor) ([]RuntimeServer, []RuntimeFailure) {
|
|
results := make([]RuntimeServer, 0, len(servers))
|
|
failures := make([]RuntimeFailure, 0)
|
|
for _, server := range servers {
|
|
if server.CLI.Skip {
|
|
continue
|
|
}
|
|
runtimeServer, err := s.DiscoverServerRuntime(ctx, server)
|
|
if err != nil {
|
|
if errors.Is(err, errCLIServerSkipped) {
|
|
continue
|
|
}
|
|
failures = append(failures, RuntimeFailure{
|
|
ServerKey: server.Key,
|
|
Err: err,
|
|
})
|
|
continue
|
|
}
|
|
results = append(results, runtimeServer)
|
|
}
|
|
return results, failures
|
|
}
|
|
|
|
func (s *Service) DiscoverDetail(ctx context.Context, server market.ServerDescriptor) (market.DetailResponse, error) {
|
|
partition := s.partition()
|
|
var fetchErr error
|
|
|
|
if detailURL := strings.TrimSpace(server.DetailLocator.DetailURL); detailURL != "" {
|
|
detail, err := s.MarketClient.FetchDetailByURL(ctx, detailURL)
|
|
if err == nil {
|
|
cacheDetailSnapshot(s.Cache, partition, server, server.DetailLocator.MCPID, detail)
|
|
s.invalidateToolsIfVersionChanged(partition, server.Key, detail)
|
|
return detail, nil
|
|
}
|
|
fetchErr = err
|
|
}
|
|
if server.DetailLocator.MCPID > 0 {
|
|
detail, err := s.MarketClient.FetchDetail(ctx, server.DetailLocator.MCPID)
|
|
if err == nil {
|
|
cacheDetailSnapshot(s.Cache, partition, server, server.DetailLocator.MCPID, detail)
|
|
s.invalidateToolsIfVersionChanged(partition, server.Key, detail)
|
|
return detail, nil
|
|
}
|
|
fetchErr = err
|
|
}
|
|
if fetchErr == nil {
|
|
fetchErr = fmt.Errorf("server %s does not expose detail locator", server.Key)
|
|
}
|
|
|
|
snapshot, _, cacheErr := s.Cache.LoadDetail(partition, server.Key)
|
|
if cacheErr != nil {
|
|
return market.DetailResponse{}, fetchErr
|
|
}
|
|
if server.DetailLocator.MCPID > 0 && snapshot.MCPID != server.DetailLocator.MCPID {
|
|
return market.DetailResponse{}, fetchErr
|
|
}
|
|
|
|
var cached market.DetailResponse
|
|
if unmarshalErr := json.Unmarshal(snapshot.Payload, &cached); unmarshalErr != nil {
|
|
return market.DetailResponse{}, fetchErr
|
|
}
|
|
return cached, nil
|
|
}
|
|
|
|
// invalidateToolsIfVersionChanged checks whether a fresh Detail API response
|
|
// contains actionVersion values that differ from those stored in the cached
|
|
// tools snapshot. If any tool's version has changed, the tools cache is
|
|
// invalidated so the next DiscoverServerRuntime call re-fetches tools/list.
|
|
func (s *Service) invalidateToolsIfVersionChanged(partition, serverKey string, detail market.DetailResponse) {
|
|
if !detail.Success || len(detail.Result.Tools) == 0 {
|
|
return
|
|
}
|
|
snapshot, _, err := s.Cache.LoadTools(partition, serverKey)
|
|
if err != nil || len(snapshot.ActionVersions) == 0 {
|
|
return
|
|
}
|
|
if cache.HasActionVersionChanged(snapshot.ActionVersions, detail.Result.Tools) {
|
|
_ = s.Cache.DeleteTools(partition, serverKey)
|
|
}
|
|
}
|
|
|
|
func (s *Service) partition() string {
|
|
return fmt.Sprintf("%s/%s", s.Tenant, s.AuthIdentity)
|
|
}
|
|
|
|
func (s *Service) CachePartition() string {
|
|
return s.partition()
|
|
}
|
|
|
|
func mergeRuntimeToolsWithDetail(tools []transport.ToolDescriptor, detail market.DetailResponse) []transport.ToolDescriptor {
|
|
if !detail.Success || len(detail.Result.Tools) == 0 || len(tools) == 0 {
|
|
return tools
|
|
}
|
|
|
|
byName := make(map[string]market.DetailTool, len(detail.Result.Tools))
|
|
for _, tool := range detail.Result.Tools {
|
|
name := strings.TrimSpace(tool.ToolName)
|
|
if name == "" {
|
|
continue
|
|
}
|
|
byName[name] = tool
|
|
}
|
|
|
|
out := make([]transport.ToolDescriptor, 0, len(tools))
|
|
for _, tool := range tools {
|
|
merged := tool
|
|
detailTool, ok := byName[strings.TrimSpace(tool.Name)]
|
|
if !ok {
|
|
out = append(out, merged)
|
|
continue
|
|
}
|
|
if title := strings.TrimSpace(detailTool.ToolTitle); title != "" {
|
|
merged.Title = title
|
|
}
|
|
if description := strings.TrimSpace(detailTool.ToolDesc); description != "" {
|
|
merged.Description = description
|
|
}
|
|
merged.Sensitive = detailTool.IsSensitive
|
|
if inputSchema := parseDetailSchema(detailTool.ToolRequest); len(inputSchema) > 0 {
|
|
merged.InputSchema = inputSchema
|
|
}
|
|
if outputSchema := parseDetailSchema(detailTool.ToolResponse); len(outputSchema) > 0 {
|
|
merged.OutputSchema = outputSchema
|
|
}
|
|
out = append(out, merged)
|
|
}
|
|
return out
|
|
}
|
|
|
|
func parseDetailSchema(raw string) map[string]any {
|
|
raw = strings.TrimSpace(raw)
|
|
if raw == "" {
|
|
return nil
|
|
}
|
|
var parsed any
|
|
if err := json.Unmarshal([]byte(raw), &parsed); err != nil {
|
|
return nil
|
|
}
|
|
object, ok := parsed.(map[string]any)
|
|
if !ok {
|
|
return nil
|
|
}
|
|
return object
|
|
}
|
|
|
|
func cacheDetailSnapshot(store *cache.Store, partition string, server market.ServerDescriptor, mcpID int, detail market.DetailResponse) {
|
|
if store == nil {
|
|
return
|
|
}
|
|
raw, err := json.Marshal(detail)
|
|
if err != nil {
|
|
return
|
|
}
|
|
snapshot := cache.DetailSnapshot{
|
|
MCPID: mcpID,
|
|
Payload: raw,
|
|
}
|
|
for _, cacheKey := range detailSnapshotKeys(server) {
|
|
_ = store.SaveDetail(partition, cacheKey, snapshot)
|
|
}
|
|
}
|
|
|
|
func detailSnapshotKeys(server market.ServerDescriptor) []string {
|
|
seen := make(map[string]struct{}, 2)
|
|
keys := make([]string, 0, 2)
|
|
for _, candidate := range []string{
|
|
strings.TrimSpace(server.Key),
|
|
strings.TrimSpace(server.CLI.ID),
|
|
} {
|
|
if candidate == "" {
|
|
continue
|
|
}
|
|
if _, ok := seen[candidate]; ok {
|
|
continue
|
|
}
|
|
seen[candidate] = struct{}{}
|
|
keys = append(keys, candidate)
|
|
}
|
|
return keys
|
|
}
|
|
|
|
func resolveTenant() string {
|
|
value := strings.TrimSpace(os.Getenv(tenantEnv))
|
|
if value == "" {
|
|
return "default"
|
|
}
|
|
return strings.ToLower(value)
|
|
}
|
|
|
|
func resolveAuthIdentity() string {
|
|
if value := strings.TrimSpace(os.Getenv(authIdentityEnv)); value != "" {
|
|
return strings.ToLower(value)
|
|
}
|
|
return "default"
|
|
}
|