diff --git a/agentscope-core/src/main/java/io/agentscope/core/tracing/Tracer.java b/agentscope-core/src/main/java/io/agentscope/core/tracing/Tracer.java index 26ae7a51f..c3d1e3c92 100644 --- a/agentscope-core/src/main/java/io/agentscope/core/tracing/Tracer.java +++ b/agentscope-core/src/main/java/io/agentscope/core/tracing/Tracer.java @@ -65,4 +65,6 @@ default List callFormat( default TResp runWithContext(ContextView reactorCtx, Supplier inner) { return inner.get(); } + + default void shutdown() {} } diff --git a/agentscope-core/src/main/java/io/agentscope/core/tracing/TracerRegistry.java b/agentscope-core/src/main/java/io/agentscope/core/tracing/TracerRegistry.java index adced638b..eb1f9c2b5 100644 --- a/agentscope-core/src/main/java/io/agentscope/core/tracing/TracerRegistry.java +++ b/agentscope-core/src/main/java/io/agentscope/core/tracing/TracerRegistry.java @@ -147,6 +147,13 @@ public static void register(Tracer tracer) { } } + public static void resetToNoop() { + Tracer previousTracer = TracerRegistry.tracer; + TracerRegistry.tracer = new NoopTracer(); + disableTracingHook(); + previousTracer.shutdown(); + } + public static Tracer get() { return tracer; } diff --git a/agentscope-core/src/test/java/io/agentscope/core/tracing/TracerRegistryTest.java b/agentscope-core/src/test/java/io/agentscope/core/tracing/TracerRegistryTest.java new file mode 100644 index 000000000..e6ab76f6e --- /dev/null +++ b/agentscope-core/src/test/java/io/agentscope/core/tracing/TracerRegistryTest.java @@ -0,0 +1,59 @@ +/* + * Copyright 2024-2026 the original author or authors. + * + * 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 io.agentscope.core.tracing; + +import static org.junit.jupiter.api.Assertions.assertEquals; +import static org.junit.jupiter.api.Assertions.assertInstanceOf; + +import java.util.concurrent.atomic.AtomicInteger; +import org.junit.jupiter.api.AfterEach; +import org.junit.jupiter.api.DisplayName; +import org.junit.jupiter.api.Test; + +@DisplayName("TracerRegistry Tests") +class TracerRegistryTest { + + @AfterEach + void tearDown() { + TracerRegistry.resetToNoop(); + } + + @Test + @DisplayName("resetToNoop() should shutdown current tracer and restore noop tracer") + void resetToNoopShouldShutdownCurrentTracer() { + CloseCountingTracer tracer = new CloseCountingTracer(); + TracerRegistry.register(tracer); + + TracerRegistry.resetToNoop(); + + assertEquals(1, tracer.shutdownCount()); + assertInstanceOf(NoopTracer.class, TracerRegistry.get()); + } + + static class CloseCountingTracer implements Tracer { + private final AtomicInteger shutdownCount = new AtomicInteger(); + + @Override + public void shutdown() { + shutdownCount.incrementAndGet(); + } + + int shutdownCount() { + return shutdownCount.get(); + } + } +} diff --git a/agentscope-extensions/agentscope-extensions-studio/src/main/java/io/agentscope/core/studio/StudioManager.java b/agentscope-extensions/agentscope-extensions-studio/src/main/java/io/agentscope/core/studio/StudioManager.java index ed3e8fd80..ea3433977 100644 --- a/agentscope-extensions/agentscope-extensions-studio/src/main/java/io/agentscope/core/studio/StudioManager.java +++ b/agentscope-extensions/agentscope-extensions-studio/src/main/java/io/agentscope/core/studio/StudioManager.java @@ -119,6 +119,7 @@ public static void shutdown() { config = null; client = null; wsClient = null; + TracerRegistry.resetToNoop(); } /** diff --git a/agentscope-extensions/agentscope-extensions-studio/src/main/java/io/agentscope/core/tracing/telemetry/TelemetryTracer.java b/agentscope-extensions/agentscope-extensions-studio/src/main/java/io/agentscope/core/tracing/telemetry/TelemetryTracer.java index d4b96d4d3..3c42a1cd2 100644 --- a/agentscope-extensions/agentscope-extensions-studio/src/main/java/io/agentscope/core/tracing/telemetry/TelemetryTracer.java +++ b/agentscope-extensions/agentscope-extensions-studio/src/main/java/io/agentscope/core/tracing/telemetry/TelemetryTracer.java @@ -67,9 +67,16 @@ public class TelemetryTracer implements Tracer { private final io.opentelemetry.api.trace.Tracer tracer; + private final SdkTracerProvider sdkTracerProvider; public TelemetryTracer(io.opentelemetry.api.trace.Tracer tracer) { + this(tracer, null); + } + + private TelemetryTracer( + io.opentelemetry.api.trace.Tracer tracer, SdkTracerProvider sdkTracerProvider) { this.tracer = tracer; + this.sdkTracerProvider = sdkTracerProvider; } @Override @@ -226,6 +233,13 @@ public TResp runWithContext(ContextView reactorCtx, Supplier inne return otelContext.wrapSupplier(inner).get(); } + @Override + public void shutdown() { + if (sdkTracerProvider != null) { + sdkTracerProvider.close(); + } + } + public static Builder builder() { return new Builder(); } @@ -294,14 +308,15 @@ public TelemetryTracer build() { exporterBuilder.addHeader(entry.getKey(), entry.getValue()); } - TracerProvider tracerProvider = + SdkTracerProvider tracerProvider = SdkTracerProvider.builder() .addSpanProcessor( BatchSpanProcessor.builder(exporterBuilder.build()).build()) .setSampler(Sampler.alwaysOn()) .build(); - return new TelemetryTracer(tracerProvider.get(INSTRUMENTATION_NAME, Version.VERSION)); + return new TelemetryTracer( + tracerProvider.get(INSTRUMENTATION_NAME, Version.VERSION), tracerProvider); } } }