| 460 | } |
| 461 | |
| 462 | func (s *Server) GetMCPServerAccessTokensBatch(ctx context.Context, in *proto.GetMCPServerAccessTokensBatchRequest) (*proto.GetMCPServerAccessTokensBatchResponse, error) { |
| 463 | if len(in.GetMcpServerConfigIds()) == 0 { |
| 464 | return &proto.GetMCPServerAccessTokensBatchResponse{}, nil |
| 465 | } |
| 466 | |
| 467 | userID, err := uuid.Parse(in.GetUserId()) |
| 468 | if err != nil { |
| 469 | return nil, xerrors.Errorf("parse user_id: %w", err) |
| 470 | } |
| 471 | |
| 472 | //nolint:gocritic // AIBridged has specific authz rules. |
| 473 | ctx = dbauthz.AsAIBridged(ctx) |
| 474 | links, err := s.store.GetExternalAuthLinksByUserID(ctx, userID) |
| 475 | if err != nil { |
| 476 | return nil, xerrors.Errorf("fetch external auth links: %w", err) |
| 477 | } |
| 478 | |
| 479 | if len(links) == 0 { |
| 480 | return &proto.GetMCPServerAccessTokensBatchResponse{}, nil |
| 481 | } |
| 482 | |
| 483 | // Ensure unique to prevent unnecessary effort. |
| 484 | ids := in.GetMcpServerConfigIds() |
| 485 | slices.Sort(ids) |
| 486 | ids = slices.Compact(ids) |
| 487 | |
| 488 | var ( |
| 489 | wg sync.WaitGroup |
| 490 | errs error |
| 491 | |
| 492 | mu sync.Mutex |
| 493 | tokens = make(map[string]string, len(ids)) |
| 494 | tokenErrs = make(map[string]string) |
| 495 | ) |
| 496 | |
| 497 | externalAuthLoop: |
| 498 | for _, id := range ids { |
| 499 | eac, ok := s.externalAuthConfigs[id] |
| 500 | if !ok { |
| 501 | mu.Lock() |
| 502 | s.logger.Warn(ctx, "no MCP server config found by given ID", slog.F("id", id)) |
| 503 | tokenErrs[id] = ErrNoMCPConfigFound.Error() |
| 504 | mu.Unlock() |
| 505 | continue |
| 506 | } |
| 507 | |
| 508 | for _, link := range links { |
| 509 | if link.ProviderID != eac.ID { |
| 510 | continue |
| 511 | } |
| 512 | |
| 513 | // Validate all configured External Auth links concurrently. |
| 514 | wg.Add(1) |
| 515 | go func() { |
| 516 | defer wg.Done() |
| 517 | |
| 518 | // TODO: timeout. |
| 519 | valid, _, validateErr := eac.ValidateToken(ctx, link.OAuthToken()) |