A365 telemetry: align with SDK 0.1.0, remove use_tenant_island_endpoint

Signed-off-by: afourniernv <afournier@nvidia.com>

Author: afourniernv

packages/nvidia_nat_a365/src/nat/plugins/a365/telemetry/a365_exporter.py

@@ -155,6 +155,8 @@ def __init__(
         resource_attributes: dict[str, str] | None = None,
         auth_provider=None,
         token_cache=None,
+        auth_ref=None,
+        builder=None,
     ):
         """Initialize the A365 exporter."""
         super().__init__(
@@ -173,8 +175,11 @@ def __init__(
         self._cluster_category = cluster_category
         self._use_s2s_endpoint = use_s2s_endpoint
         self._suppress_invoke_agent_input = suppress_invoke_agent_input
-        self._auth_provider = auth_provider  # For proactive token refresh
-        self._token_cache = token_cache  # For updating cached tokens
+        self._auth_provider = auth_provider
+        self._token_cache = token_cache
+        self._auth_ref = auth_ref
+        self._builder = builder
+        self._auth_resolve_lock = asyncio.Lock()
 
         # SDK requires token_resolver to be non-None, so if None is passed, SDK will raise ValueError
         self._a365_exporter = Agent365Exporter(
@@ -188,6 +193,52 @@ def __init__(
             f"tenant_id={tenant_id}, cluster={cluster_category}"
         )
 
+    async def _resolve_auth_once(self) -> None:
+        """Resolve auth provider and fill token cache on first export (lazy).
+
+        Telemetry is built in __aenter__ before auth exists; by first export,
+        populate_builder has run so we can resolve here. Keeps core unchanged.
+        """
+        if self._auth_provider is not None or self._auth_ref is None or self._builder is None:
+            return
+        if self._token_cache is None:
+            return
+        async with self._auth_resolve_lock:
+            if self._auth_provider is not None:
+                return
+            try:
+                auth_provider = await self._builder.get_auth_provider(self._auth_ref)
+                from nat.builder.context import Context
+                user_id = Context.get().user_id
+                auth_result = await auth_provider.authenticate(user_id=user_id)
+                if not auth_result.credentials:
+                    raise A365AuthenticationError("No credentials available from auth provider")
+                from nat.data_models.authentication import BearerTokenCred, HeaderCred
+                from nat.authentication.interfaces import AUTHORIZATION_HEADER
+                token = None
+                for cred in auth_result.credentials:
+                    if isinstance(cred, BearerTokenCred):
+                        token = cred.token.get_secret_value()
+                        break
+                    if isinstance(cred, HeaderCred) and cred.name == AUTHORIZATION_HEADER:
+                        hv = cred.value.get_secret_value()
+                        token = hv[7:] if hv.startswith("Bearer ") else hv
+                        break
+                if token is None:
+                    raise A365AuthenticationError(
+                        f"No bearer token in credentials. "
+                        f"Types: {[type(c).__name__ for c in auth_result.credentials]}"
+                    )
+                self._token_cache.update_token(token, auth_result.token_expires_at)
+                self._auth_provider = auth_provider
+            except Exception as e:
+                logger.error(
+                    f"Failed to resolve auth on first export (agent_id={self._agent_id}, "
+                    f"tenant_id={self._tenant_id}): {e}",
+                    exc_info=True,
+                )
+                raise
+
     async def _refresh_token_if_needed(self) -> None:
         """Refresh token proactively if it's expiring soon.
 
@@ -261,6 +312,7 @@ async def export_otel_spans(self, spans: list[OtelSpan]) -> None:
         if not spans:
             return
 
+        await self._resolve_auth_once()
         await self._refresh_token_if_needed()
 
         try:

packages/nvidia_nat_a365/src/nat/plugins/a365/telemetry/register.py

@@ -76,21 +76,32 @@ def __init__(self, token: str, expires_at: datetime | None):
         self._token = token
         self._expires_at = expires_at
 
+    def _expires_at_utc(self) -> datetime | None:
+        """Return expiration as timezone-aware UTC for comparison.
+
+        If naive, assume local time (e.g. from provider using datetime.now() + delta).
+        """
+        if self._expires_at is None:
+            return None
+        if self._expires_at.tzinfo is not None:
+            return self._expires_at
+        local_tz = datetime.now().astimezone().tzinfo
+        return self._expires_at.replace(tzinfo=local_tz).astimezone(timezone.utc)
+
     def get_token(self) -> str | None:
         """Get cached token if still valid (with 5 minute buffer).
 
         Returns:
             Token string if valid, None if expired or expiring soon
         """
         with self._lock:
-            if self._expires_at:
+            expires_utc = self._expires_at_utc()
+            if expires_utc is not None:
                 buffer_time = datetime.now(timezone.utc) + timedelta(minutes=5)
-                if self._expires_at > buffer_time:
+                if expires_utc > buffer_time:
                     return self._token
                 return None
-            else:
-                # No expiration info, assume token is valid
-                return self._token
+            return self._token
 
     def update_token(self, token: str, expires_at: datetime | None) -> None:
         """Update cached token.
@@ -113,10 +124,11 @@ def is_expiring_soon(self, buffer_minutes: int = 5) -> bool:
             True if token expires within buffer time, False otherwise
         """
         with self._lock:
-            if self._expires_at is None:
+            expires_utc = self._expires_at_utc()
+            if expires_utc is None:
                 return False
             buffer_time = datetime.now(timezone.utc) + timedelta(minutes=buffer_minutes)
-            return self._expires_at <= buffer_time
+            return expires_utc <= buffer_time
 
 
 async def _create_token_resolver_from_auth_ref(
@@ -195,25 +207,35 @@ async def a365_telemetry_exporter(config: A365TelemetryExporter, builder: Builde
 
     Integrates A365's Agent365Exporter with NAT's telemetry system to send
     OpenTelemetry spans to Microsoft Agent 365 backend endpoints.
+
+    Auth is resolved lazily on first export (not at build time) because
+    telemetry exporters are built in WorkflowBuilder.__aenter__ before
+    auth providers exist. This keeps core unchanged; the plugin handles
+    the timing constraint locally.
     """
     from nat.plugins.a365.telemetry.a365_exporter import A365OtelExporter
 
-    token_resolver_callable, auth_provider, token_cache = await _create_token_resolver_from_auth_ref(
-        config.token_resolver, builder
-    )
+    # Defer auth: do not call get_auth_provider here (not available yet in __aenter__).
+    token_cache = _TokenCache(None, None)
+
+    def token_resolver(agent_id: str, tenant_id: str) -> str | None:
+        """Sync callable for SDK; returns cached token (filled on first export)."""
+        return token_cache.get_token()
 
     logger.info(
         f"A365 telemetry exporter initialized for agent_id={config.agent_id}, "
         f"tenant_id={config.tenant_id}, cluster={config.cluster_category}, "
-        f"token_resolver=configured (auth_provider='{config.token_resolver}')"
+        f"token_resolver=deferred (auth resolved on first export)"
     )
 
     exporter = A365OtelExporter(
         agent_id=config.agent_id,
         tenant_id=config.tenant_id,
-        token_resolver=token_resolver_callable,
-        auth_provider=auth_provider,  # Pass auth provider for proactive refresh
-        token_cache=token_cache,  # Pass token cache for updating tokens
+        token_resolver=token_resolver,
+        auth_provider=None,
+        token_cache=token_cache,
+        auth_ref=config.token_resolver,
+        builder=builder,
         cluster_category=config.cluster_category,
         use_s2s_endpoint=config.use_s2s_endpoint,
         suppress_invoke_agent_input=config.suppress_invoke_agent_input,

packages/nvidia_nat_a365/tests/telemetry/test_registration_integration.py

@@ -76,16 +76,32 @@ async def test_registration_creates_exporter_successfully(
 
             with patch("nat.plugins.a365.telemetry.a365_exporter.Agent365Exporter"):
                 async with a365_telemetry_exporter(config, mock_builder) as exporter:
-                    # Verify exporter was created
+                    # Lazy init: auth not resolved at build time (exporter built in __aenter__ before auth exists)
                     assert exporter is not None
                     assert exporter._agent_id == "test-agent-123"
                     assert exporter._tenant_id == "test-tenant-456"
                     assert exporter._token_resolver is not None
-                    assert exporter._auth_provider == mock_auth_provider
+                    assert exporter._auth_provider is None
+                    assert exporter._auth_ref == config.token_resolver
+                    assert exporter._builder is mock_builder
                     assert exporter._token_cache is not None
-
-                    # Verify auth provider was resolved
+                    mock_builder.get_auth_provider.assert_not_called()
+                    # Resolve on first export (minimal span-like object for conversion)
+                    mock_span = Mock()
+                    mock_span.attributes = {}
+                    mock_span.events = []
+                    mock_span.links = []
+                    mock_span.name = "test"
+                    mock_span.kind = 1
+                    mock_span.start_time = 0
+                    mock_span.end_time = 0
+                    mock_span.status = None
+                    mock_span.instrumentation_scope = None
+                    mock_span.resource = None
+                    mock_span.parent = None
+                    await exporter.export_otel_spans([mock_span])
                     mock_builder.get_auth_provider.assert_called_once_with(config.token_resolver)
+                    assert exporter._auth_provider is mock_auth_provider
 
     @pytest.mark.asyncio
     async def test_registration_passes_all_config_to_exporter(
@@ -124,42 +140,48 @@ async def test_registration_passes_all_config_to_exporter(
     async def test_registration_handles_auth_provider_resolution_failure(
         self, config, mock_builder
     ):
-        """Test that registration handles auth provider resolution failure."""
+        """Test that auth provider resolution failure is raised on first export (lazy)."""
         mock_builder.get_auth_provider.side_effect = ValueError("Auth provider not found")
 
-        with pytest.raises(ValueError, match="Auth provider not found"):
-            async with a365_telemetry_exporter(config, mock_builder):
-                pass
+        with patch("nat.plugins.a365.telemetry.a365_exporter.Agent365Exporter"):
+            async with a365_telemetry_exporter(config, mock_builder) as exporter:
+                mock_span = Mock()
+                with pytest.raises(ValueError, match="Auth provider not found"):
+                    await exporter.export_otel_spans([mock_span])
 
     @pytest.mark.asyncio
     async def test_registration_handles_authentication_failure(
         self, config, mock_builder, mock_auth_provider
     ):
-        """Test that registration handles authentication failure."""
+        """Test that authentication failure is raised on first export (lazy)."""
         mock_auth_provider.authenticate.side_effect = A365AuthenticationError("Auth failed")
 
         with patch("nat.builder.context.Context") as mock_context_class:
             mock_context_class.get.return_value.user_id = "test_user"
 
-            with pytest.raises(A365AuthenticationError, match="Auth failed"):
-                async with a365_telemetry_exporter(config, mock_builder):
-                    pass
+            with patch("nat.plugins.a365.telemetry.a365_exporter.Agent365Exporter"):
+                async with a365_telemetry_exporter(config, mock_builder) as exporter:
+                    mock_span = Mock()
+                    with pytest.raises(A365AuthenticationError, match="Auth failed"):
+                        await exporter.export_otel_spans([mock_span])
 
     @pytest.mark.asyncio
     async def test_registration_handles_no_credentials(
         self, config, mock_builder, mock_auth_provider
     ):
-        """Test that registration handles case when auth provider returns no credentials."""
+        """Test that no credentials is raised on first export (lazy)."""
         auth_result = Mock(spec=AuthResult)
         auth_result.credentials = []
         mock_auth_provider.authenticate.return_value = auth_result
 
         with patch("nat.builder.context.Context") as mock_context_class:
             mock_context_class.get.return_value.user_id = "test_user"
 
-            with pytest.raises(A365AuthenticationError, match="No credentials available"):
-                async with a365_telemetry_exporter(config, mock_builder):
-                    pass
+            with patch("nat.plugins.a365.telemetry.a365_exporter.Agent365Exporter"):
+                async with a365_telemetry_exporter(config, mock_builder) as exporter:
+                    mock_span = Mock()
+                    with pytest.raises(A365AuthenticationError, match="No credentials available"):
+                        await exporter.export_otel_spans([mock_span])
 
     @pytest.mark.asyncio
     async def test_registration_is_context_manager(self, config, mock_builder, mock_auth_provider):