diff --git a/agent-memory-client/agent-memory-client-java/build.gradle.kts b/agent-memory-client/agent-memory-client-java/build.gradle.kts index b02e86e8..c7d3da29 100644 --- a/agent-memory-client/agent-memory-client-java/build.gradle.kts +++ b/agent-memory-client/agent-memory-client-java/build.gradle.kts @@ -36,7 +36,8 @@ repositories { dependencies { // HTTP Client - implementation("com.squareup.okhttp3:okhttp:4.12.0") + // api so that `Builder.httpClient` callers can avoid version skew + api("com.squareup.okhttp3:okhttp:4.12.0") // JSON Processing implementation("com.fasterxml.jackson.core:jackson-databind:2.16.1") diff --git a/agent-memory-client/agent-memory-client-java/src/main/java/com/redis/agentmemory/MemoryAPIClient.java b/agent-memory-client/agent-memory-client-java/src/main/java/com/redis/agentmemory/MemoryAPIClient.java index fc0a2997..e22a2b62 100644 --- a/agent-memory-client/agent-memory-client-java/src/main/java/com/redis/agentmemory/MemoryAPIClient.java +++ b/agent-memory-client/agent-memory-client-java/src/main/java/com/redis/agentmemory/MemoryAPIClient.java @@ -36,6 +36,7 @@ public class MemoryAPIClient implements AutoCloseable { private final String defaultModelName; private final Integer defaultContextWindowMax; private final OkHttpClient httpClient; + private final boolean ownsHttpClient; private final ObjectMapper objectMapper; // Service instances @@ -53,10 +54,17 @@ private MemoryAPIClient(Builder builder) { this.defaultModelName = builder.defaultModelName; this.defaultContextWindowMax = builder.defaultContextWindowMax; - this.httpClient = new OkHttpClient.Builder() - .connectTimeout((long) timeout, TimeUnit.SECONDS) - .readTimeout((long) timeout, TimeUnit.SECONDS) - .writeTimeout((long) timeout, TimeUnit.SECONDS) + OkHttpClient.Builder httpClientBuilder = builder.httpClient != null + ? builder.httpClient.newBuilder() + : new OkHttpClient.Builder() + .connectTimeout((long) timeout, TimeUnit.SECONDS) + .readTimeout((long) timeout, TimeUnit.SECONDS) + .writeTimeout((long) timeout, TimeUnit.SECONDS); + + // when the caller provides a client, we don't own it's lifecycle + this.ownsHttpClient = builder.httpClient == null; + + this.httpClient = httpClientBuilder .addInterceptor(chain -> { Request original = chain.request(); Request request = original.newBuilder() @@ -161,8 +169,11 @@ public TaskService tasks() { @Override public void close() { - httpClient.dispatcher().executorService().shutdown(); - httpClient.connectionPool().evictAll(); + // when the caller provides a client, we don't own it's lifecycle + if (ownsHttpClient) { + httpClient.dispatcher().executorService().shutdown(); + httpClient.connectionPool().evictAll(); + } } /** @@ -355,6 +366,7 @@ public static class Builder { private String defaultNamespace = null; private String defaultModelName = null; private Integer defaultContextWindowMax = null; + private OkHttpClient httpClient = null; private Builder(@NotNull String baseUrl) { this.baseUrl = baseUrl; @@ -400,6 +412,19 @@ public Builder defaultContextWindowMax(@Nullable Integer defaultContextWindowMax return this; } + /** + * Sets an {@link OkHttpClient} to use for all requests. + * + *

+ * When provided, the {@link #timeout(double)} settings is ignored. + * @param httpClient the OkHttpClient to use, or null to use the default + * @return this builder + */ + public Builder httpClient(@Nullable OkHttpClient httpClient) { + this.httpClient = httpClient; + return this; + } + /** * Builds the MemoryAPIClient instance. * @return a new MemoryAPIClient diff --git a/agent-memory-client/agent-memory-client-java/src/test/java/com/redis/agentmemory/MemoryAPIClientTest.java b/agent-memory-client/agent-memory-client-java/src/test/java/com/redis/agentmemory/MemoryAPIClientTest.java index 220d0600..9a097401 100644 --- a/agent-memory-client/agent-memory-client-java/src/test/java/com/redis/agentmemory/MemoryAPIClientTest.java +++ b/agent-memory-client/agent-memory-client-java/src/test/java/com/redis/agentmemory/MemoryAPIClientTest.java @@ -6,6 +6,7 @@ import com.redis.agentmemory.models.common.AckResponse; import com.redis.agentmemory.models.longtermemory.*; import com.redis.agentmemory.models.workingmemory.*; +import okhttp3.OkHttpClient; import okhttp3.mockwebserver.MockResponse; import okhttp3.mockwebserver.MockWebServer; import okhttp3.mockwebserver.RecordedRequest; @@ -404,5 +405,61 @@ void testPromoteWorkingMemoriesToLongTerm_NoMemories() throws Exception { assertEquals(1, mockServer.getRequestCount()); // Only GET, no CREATE } + @Test + void testCustomHttpClientReachesServer() throws Exception { + // caller-supplied client + OkHttpClient customClient = new OkHttpClient.Builder() + .addInterceptor(chain -> chain.proceed( + chain.request().newBuilder() + .header("X-Custom-Header", "custom-header-value") + .build())) + .build(); + + try (MemoryAPIClient customizedClient = MemoryAPIClient.builder(mockServer.url("/").toString()) + .httpClient(customClient) + .build()) { + + mockServer.enqueue(new MockResponse() + .setBody("{}") + .addHeader("Content-Type", "application/json")); + + customizedClient.health().healthCheck(); + + RecordedRequest request = mockServer.takeRequest(); + // provided client headers + assertEquals("custom-header-value", request.getHeader("X-Custom-Header")); + // normal headers + assertTrue(request.getHeader("User-Agent").startsWith("agent-memory-client-java/")); + assertNotNull(request.getHeader("X-Client-Version")); + } + } + + @Test + void testDefaultHttpClientHeaders() throws Exception { + mockServer.enqueue(new MockResponse() + .setBody("{}") + .addHeader("Content-Type", "application/json")); + + client.health().healthCheck(); + + RecordedRequest request = mockServer.takeRequest(); + assertTrue(request.getHeader("User-Agent").startsWith("agent-memory-client-java/")); + assertNotNull(request.getHeader("X-Client-Version")); + } + + @Test + void testCloseDoesNotShutDownCallerSuppliedClient() throws Exception { + OkHttpClient customClient = new OkHttpClient.Builder().build(); + + MemoryAPIClient customizedClient = MemoryAPIClient.builder(mockServer.url("/").toString()) + .httpClient(customClient) + .build(); + + customizedClient.close(); + + assertFalse(customClient.dispatcher().executorService().isShutdown(), + "caller-supplied client's executor must survive MemoryAPIClient.close()"); + } + // ===== Tests for Phase 2: Convenience Overloads ===== }