Java SDK

The 2Signal Java SDK instruments your Java and Kotlin AI agents with automatic tracing, token counting, and cost tracking. Java 11+. Uses ThreadLocal for context propagation and daemon threads for background export.

Installation

Maven

<dependency>
  <groupId>com.twosignal</groupId>
  <artifactId>twosignal-sdk</artifactId>
  <version>0.1.0</version>
</dependency>

Gradle

implementation 'com.twosignal:twosignal-sdk:0.1.0'

Quick Start

import com.twosignal.sdk.TwoSignal;
import com.twosignal.sdk.TwoSignalOptions;
import com.twosignal.sdk.Models.SpanType;

public class MyAgent {
    public static void main(String[] args) {
        TwoSignal ts = new TwoSignal(TwoSignalOptions.builder()
            .apiKey("ts_...")  // or set TWOSIGNAL_API_KEY env var
            .build());

        String result = ts.span("support-agent", SpanType.AGENT)
            .input("What is your refund policy?")
            .run(span -> {
                // your agent logic here
                String answer = "Our refund policy is...";
                span.setOutput(answer);
                return answer;
            });

        System.out.println(result);
        ts.shutdown();
    }
}

Initialization

TwoSignal ts = new TwoSignal(TwoSignalOptions.builder()
    .apiKey("ts_...")          // or TWOSIGNAL_API_KEY env var
    .baseUrl("https://...")    // or TWOSIGNAL_BASE_URL env var
    .enabled(true)             // set false to disable all tracing
    .flushIntervalMs(1000)     // milliseconds between batch flushes
    .maxBatchSize(100)         // max events per batch
    .sessionId("session-123")  // optional session grouping
    .build());

// the constructor sets the global instance automatically
TwoSignal instance = TwoSignal.getInstance();

The SDK is a singleton. It starts a daemon thread that batches and flushes events. Call shutdown() before your process exits to flush remaining events.

Configuration

Creating Spans

// With return value
String result = ts.span("classify-intent", SpanType.TOOL)
    .input(userQuery)
    .metadata(Map.of("model", "gpt-4o"))
    .run(span -> {
        String classification = classifier.classify(userQuery);
        span.setOutput(classification);
        return classification;
    });

// Without return value (void)
ts.span("log-event", SpanType.CUSTOM)
    .input(eventData)
    .run(span -> {
        eventLogger.log(eventData);
        span.setOutput("logged");
    });

Fields you can set on the span inside run():

If your function throws an exception, the span is automatically marked as ERROR with the exception message recorded. The exception is re-thrown to the caller.

Context Propagation

The SDK uses ThreadLocal to propagate trace context. Nested spans automatically become children of their parent:

ts.span("outer-agent", SpanType.AGENT).run(outerSpan -> {
    // this span automatically becomes a child of "outer-agent"
    ts.span("inner-tool", SpanType.TOOL).run(innerSpan -> {
        // nested correctly — same trace ID, parent set automatically
        innerSpan.setOutput("done");
    });
});

Span Types

OpenAI Wrapper

If you use the openai-java SDK, the wrapper auto-traces chat completions:

import com.twosignal.sdk.wrappers.OpenAIWrapper;

// Wrap your OpenAI client
OpenAIWrapper wrapper = new OpenAIWrapper(openAIClient);

// This automatically creates an LLM span with model, usage, and cost
ChatCompletion response = wrapper.tracedChatCompletion(
    ChatCompletionCreateParams.builder()
        .model("gpt-4o")
        .addMessage(...)
        .build()
);

Cost Calculation

import com.twosignal.sdk.CostCalculator;

Double cost = CostCalculator.calculateCost("gpt-4o", 100, 50);
// returns null if model not in pricing table

Supports OpenAI, Anthropic, Google Gemini, Mistral, Cohere, and Groq models.

Manual LLM Span Example

ts.span("openai-chat", SpanType.LLM)
    .input(messages)
    .run(span -> {
        ChatCompletion resp = openai.chatCompletion(request);

        span.setModel("gpt-4o");
        span.setOutput(resp.choices().get(0).message().content());
        span.setUsage(new Models.Usage(
            resp.usage().promptTokens(),
            resp.usage().completionTokens(),
            resp.usage().totalTokens()
        ));
        Double cost = CostCalculator.calculateCost("gpt-4o",
            resp.usage().promptTokens(),
            resp.usage().completionTokens());
        if (cost != null) span.setCost(cost);
    });

Lifecycle

// force-flush all pending events
ts.flush();

// graceful shutdown — flushes and stops the daemon thread (5s timeout)
ts.shutdown();

Always call shutdown() before your application exits to ensure all events are flushed. In Spring Boot, register a @PreDestroy method. In plain Java, use a shutdown hook.

Thread Safety

The SDK is fully thread-safe. The singleton instance, exporter queue, and context propagation all use proper synchronization. You can safely create spans from any thread.

Note: ThreadLocal context does not propagate across thread boundaries. If you spawn new threads (e.g. ExecutorService), context will not automatically carry over. For cross-thread tracing, manually pass the trace ID.