Skip to content

Commit 61f3523

Browse files
JqRrtJqRrt
committed
rebuild sdk
1 parent 6b079a0 commit 61f3523

67 files changed

Lines changed: 2974 additions & 1692 deletions

File tree

Some content is hidden

Large Commits have some content hidden by default. Use the searchbox below for content that may be hidden.

‎cozeloop-core/pom.xml‎

Lines changed: 5 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -41,6 +41,11 @@
4141
<version>1.23.1-alpha</version>
4242
</dependency>
4343

44+
<dependency>
45+
<groupId>io.opentelemetry</groupId>
46+
<artifactId>opentelemetry-exporter-otlp</artifactId>
47+
</dependency>
48+
4449
<!-- HTTP Client -->
4550
<dependency>
4651
<groupId>com.squareup.okhttp3</groupId>

‎cozeloop-core/src/main/java/com/coze/loop/auth/JWTOAuthAuth.java‎

Lines changed: 87 additions & 15 deletions
Original file line numberDiff line numberDiff line change
@@ -2,9 +2,15 @@
22

33
import com.coze.loop.exception.AuthException;
44
import com.coze.loop.exception.ErrorCode;
5+
import com.coze.loop.http.HttpClient;
6+
import com.coze.loop.internal.JsonUtils;
57
import com.coze.loop.internal.ValidationUtils;
8+
import com.fasterxml.jackson.annotation.JsonProperty;
69
import io.jsonwebtoken.Jwts;
710
import io.jsonwebtoken.SignatureAlgorithm;
11+
import okhttp3.MediaType;
12+
import okhttp3.Request;
13+
import okhttp3.RequestBody;
814
import org.slf4j.Logger;
915
import org.slf4j.LoggerFactory;
1016

@@ -13,6 +19,9 @@
1319
import java.security.spec.PKCS8EncodedKeySpec;
1420
import java.util.Base64;
1521
import java.util.Date;
22+
import java.util.HashMap;
23+
import java.util.Map;
24+
import java.util.UUID;
1625
import java.util.concurrent.locks.ReentrantReadWriteLock;
1726

1827
/**
@@ -22,36 +31,69 @@
2231
public class JWTOAuthAuth implements Auth {
2332
private static final Logger logger = LoggerFactory.getLogger(JWTOAuthAuth.class);
2433
private static final String AUTH_TYPE = "Bearer";
25-
private static final long TOKEN_EXPIRY_MINUTES = 55; // 55 minutes, JWT typically expires in 1 hour
26-
private static final long REFRESH_BUFFER_MINUTES = 5; // Refresh 5 minutes before expiry
34+
private static final String GRANT_TYPE_JWT = "urn:ietf:params:oauth:grant-type:jwt-bearer";
35+
private static final String TOKEN_PATH = "/api/permission/oauth2/token";
36+
private static final long REFRESH_BUFFER_MS = 5 * 60 * 1000; // 5 minutes
37+
private static final long DEFAULT_JWT_TTL_MS = 15 * 60 * 1000; // 15 minutes
2738

2839
private final String clientId;
2940
private final PrivateKey privateKey;
3041
private final String publicKeyId;
42+
private String baseUrl;
43+
private HttpClient httpClient;
3144

3245
private volatile String currentToken;
3346
private volatile long tokenExpiryTime;
3447

3548
private final ReentrantReadWriteLock lock = new ReentrantReadWriteLock();
3649

3750
/**
38-
* Create a JWTOAuthAuth instance.
51+
* Create a JWTOAuthAuth instance with default base URL.
3952
*
4053
* @param clientId the client ID
4154
* @param privateKeyPem the private key in PEM format (base64 encoded PKCS8)
4255
* @param publicKeyId the public key ID
4356
*/
4457
public JWTOAuthAuth(String clientId, String privateKeyPem, String publicKeyId) {
58+
this(clientId, privateKeyPem, publicKeyId, "https://api.coze.cn");
59+
}
60+
61+
/**
62+
* Create a JWTOAuthAuth instance with custom base URL.
63+
*
64+
* @param clientId the client ID
65+
* @param privateKeyPem the private key in PEM format (base64 encoded PKCS8)
66+
* @param publicKeyId the public key ID
67+
* @param baseUrl the base URL for token refresh
68+
*/
69+
public JWTOAuthAuth(String clientId, String privateKeyPem, String publicKeyId, String baseUrl) {
4570
ValidationUtils.requireNonEmpty(clientId, "clientId");
4671
ValidationUtils.requireNonEmpty(privateKeyPem, "privateKeyPem");
4772
ValidationUtils.requireNonEmpty(publicKeyId, "publicKeyId");
4873

4974
this.clientId = clientId;
5075
this.publicKeyId = publicKeyId;
5176
this.privateKey = parsePrivateKey(privateKeyPem);
52-
53-
// Generate initial token
54-
refreshToken();
77+
this.baseUrl = baseUrl != null ? baseUrl : "https://api.coze.cn";
78+
this.httpClient = new HttpClient(null);
79+
}
80+
81+
/**
82+
* Set the base URL for token refresh.
83+
*
84+
* @param baseUrl the base URL
85+
*/
86+
public void setBaseUrl(String baseUrl) {
87+
this.baseUrl = baseUrl;
88+
}
89+
90+
/**
91+
* Set the HttpClient for token refresh.
92+
*
93+
* @param httpClient the HttpClient
94+
*/
95+
public void setHttpClient(HttpClient httpClient) {
96+
this.httpClient = httpClient;
5597
}
5698

5799
@Override
@@ -91,35 +133,65 @@ public String getType() {
91133
* Check if token should be refreshed.
92134
*/
93135
private boolean shouldRefreshToken(long currentTime) {
94-
return currentTime >= (tokenExpiryTime - REFRESH_BUFFER_MINUTES * 60 * 1000);
136+
return currentToken == null || currentTime >= (tokenExpiryTime - REFRESH_BUFFER_MS);
95137
}
96138

97139
/**
98-
* Refresh the JWT token.
140+
* Refresh the OAuth token by exchanging a JWT for an access token.
99141
* Must be called with write lock held.
100142
*/
101143
private void refreshToken() {
102144
try {
103145
long currentTime = System.currentTimeMillis();
104-
long expiryTime = currentTime + TOKEN_EXPIRY_MINUTES * 60 * 1000;
146+
String host = new java.net.URL(baseUrl).getHost();
105147

148+
// 1. Generate JWT
106149
String jwt = Jwts.builder()
107150
.setIssuer(clientId)
108-
.setAudience("coze")
151+
.setAudience(host)
109152
.setIssuedAt(new Date(currentTime))
110-
.setExpiration(new Date(expiryTime))
153+
.setExpiration(new Date(currentTime + DEFAULT_JWT_TTL_MS))
111154
.setHeaderParam("kid", publicKeyId)
155+
.setHeaderParam("typ", "JWT")
156+
.setHeaderParam("alg", "RS256")
157+
.setId(UUID.randomUUID().toString())
112158
.signWith(privateKey, SignatureAlgorithm.RS256)
113159
.compact();
114160

115-
this.currentToken = jwt;
116-
this.tokenExpiryTime = expiryTime;
161+
// 2. Exchange JWT for Access Token using HttpClient
162+
Map<String, Object> body = new HashMap<>();
163+
body.put("client_id", clientId);
164+
body.put("grant_type", GRANT_TYPE_JWT);
165+
166+
Map<String, String> headers = new HashMap<>();
167+
headers.put("Authorization", "Bearer " + jwt);
168+
169+
String responseJson = httpClient.post(baseUrl + TOKEN_PATH, body, headers);
170+
TokenResponse resp = JsonUtils.fromJson(responseJson, TokenResponse.class);
117171

118-
logger.debug("JWT token refreshed, expires at: {}", new Date(expiryTime));
172+
if (resp == null || resp.accessToken == null) {
173+
throw new AuthException(ErrorCode.AUTH_FAILED, "Invalid response from token endpoint");
174+
}
175+
176+
this.currentToken = resp.accessToken;
177+
this.tokenExpiryTime = System.currentTimeMillis() + resp.expiresIn * 1000;
178+
179+
logger.debug("OAuth token refreshed, expires at: {}", new Date(this.tokenExpiryTime));
119180
} catch (Exception e) {
120-
throw new AuthException(ErrorCode.AUTH_FAILED, "Failed to generate JWT token", e);
181+
throw new AuthException(ErrorCode.AUTH_FAILED, "Failed to refresh OAuth token", e);
121182
}
122183
}
184+
185+
private static class TokenResponse {
186+
@JsonProperty("access_token")
187+
private String accessToken;
188+
189+
@JsonProperty("expires_in")
190+
private long expiresIn;
191+
192+
@JsonProperty("refresh_token")
193+
private String refreshToken;
194+
}
123195

124196
/**
125197
* Parse private key from PEM format.

‎cozeloop-core/src/main/java/com/coze/loop/client/CozeLoopClient.java‎

Lines changed: 41 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -8,6 +8,7 @@
88
import com.coze.loop.stream.StreamReader;
99
import com.coze.loop.trace.CozeLoopSpan;
1010
import io.opentelemetry.api.trace.Tracer;
11+
import io.opentelemetry.context.Context;
1112

1213
import java.util.List;
1314
import java.util.Map;
@@ -27,7 +28,7 @@ public interface CozeLoopClient extends AutoCloseable {
2728
* @return the span wrapper
2829
*/
2930
CozeLoopSpan startSpan(String name);
30-
31+
3132
/**
3233
* Start a new span with specified type.
3334
*
@@ -36,6 +37,40 @@ public interface CozeLoopClient extends AutoCloseable {
3637
* @return the span wrapper
3738
*/
3839
CozeLoopSpan startSpan(String name, String spanType);
40+
41+
CozeLoopSpan startSpan(String name, String spanType, String scene);
42+
43+
/**
44+
* Start a new span with specified type.
45+
*
46+
* @param name the span name
47+
* @param spanType the span type (e.g., "llm", "tool", "custom")
48+
* @return the span wrapper
49+
*/
50+
CozeLoopSpan startSpan(String name, String spanType, CozeLoopSpan parentSpan);
51+
52+
CozeLoopSpan startSpan(String name, String spanType, CozeLoopSpan parentSpan, String scene);
53+
54+
/**
55+
* Start a new span with specified type and parent context.
56+
* Useful for cross-service tracing where you have an extracted Context.
57+
*
58+
* @param name the span name
59+
* @param spanType the span type
60+
* @param parentContext the parent context
61+
* @return the span wrapper
62+
*/
63+
CozeLoopSpan startSpan(String name, String spanType, Context parentContext);
64+
65+
CozeLoopSpan startSpan(String name, String spanType, Context parentContext, String scene);
66+
67+
/**
68+
* Extract context from headers for cross-service propagation.
69+
*
70+
* @param headers the map of headers (e.g. from an incoming HTTP request)
71+
* @return the extracted context
72+
*/
73+
Context extractContext(Map<String, String> headers);
3974

4075
/**
4176
* Get the underlying OpenTelemetry Tracer.
@@ -103,6 +138,11 @@ public interface CozeLoopClient extends AutoCloseable {
103138
* @return workspace ID
104139
*/
105140
String getWorkspaceId();
141+
142+
/**
143+
* Flush all pending spans to the backend.
144+
*/
145+
void flush();
106146

107147
/**
108148
* Shutdown the client and release resources.

‎cozeloop-core/src/main/java/com/coze/loop/client/CozeLoopClientBuilder.java‎

Lines changed: 34 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -7,6 +7,7 @@
77
import com.coze.loop.exception.CozeLoopException;
88
import com.coze.loop.exception.ErrorCode;
99
import com.coze.loop.http.HttpClient;
10+
import com.coze.loop.internal.ConfigUtils;
1011
import com.coze.loop.internal.ValidationUtils;
1112
import com.coze.loop.prompt.PromptProvider;
1213
import com.coze.loop.trace.CozeLoopTracerProvider;
@@ -20,6 +21,27 @@ public class CozeLoopClientBuilder {
2021

2122
public CozeLoopClientBuilder() {
2223
this.config = CozeLoopConfig.builder().build();
24+
loadByProperties();
25+
}
26+
27+
private void loadByProperties() {
28+
// Load workspace ID
29+
String workspaceId = ConfigUtils.get("cozeloop.workspace-id");
30+
if (workspaceId != null) {
31+
this.config.setWorkspaceId(workspaceId);
32+
}
33+
34+
// Load service name
35+
String serviceName = ConfigUtils.get("cozeloop.service-name");
36+
if (serviceName != null) {
37+
this.config.setServiceName(serviceName);
38+
}
39+
40+
// Load auth token
41+
String token = ConfigUtils.get("cozeloop.auth.token");
42+
if (token != null) {
43+
this.tokenAuth(token);
44+
}
2345
}
2446

2547
/**
@@ -112,19 +134,25 @@ public CozeLoopClient build() {
112134
ValidationUtils.requireNonNull(auth, "auth");
113135

114136
try {
115-
// Create HTTP client
116-
HttpClient httpClient = new HttpClient(auth, config.getHttpConfig());
117-
118-
// Create TracerProvider
137+
// Create TracerProvider using OTLP/HTTP exporter
119138
CozeLoopTracerProvider tracerProvider = CozeLoopTracerProvider.create(
120-
httpClient,
139+
auth,
121140
config.getSpanEndpoint(),
122-
config.getFileEndpoint(),
123141
config.getWorkspaceId(),
124142
config.getServiceName(),
125143
config.getTraceConfig()
126144
);
127145

146+
// Create HTTP client with trace propagators
147+
HttpClient httpClient = new HttpClient(auth, config.getHttpConfig(), tracerProvider.getPropagators());
148+
149+
// Configure JWT auth if used
150+
if (auth instanceof JWTOAuthAuth) {
151+
JWTOAuthAuth jwtAuth = (JWTOAuthAuth) auth;
152+
jwtAuth.setBaseUrl(config.getBaseUrl());
153+
jwtAuth.setHttpClient(new HttpClient(null, config.getHttpConfig(), tracerProvider.getPropagators()));
154+
}
155+
128156
// Create PromptProvider
129157
PromptProvider promptProvider = new PromptProvider(
130158
httpClient,

0 commit comments

Comments
 (0)