Skip to content

fix: Preserve tokens and TTL in working memory update methods #181

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
raphaeldelio merged 1 commit into from
Mar 17, 2026
+493 −6
Merged

fix: Preserve tokens and TTL in working memory update methods #181

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
80 changes: 74 additions & 6 deletions ...memory-client-java/src/main/java/com/redis/agentmemory/services/WorkingMemoryService.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -365,10 +365,12 @@ public WorkingMemoryResponse setWorkingMemoryData(
.data(data)
.context(existing.getContext())
.userId(existing.getUserId())
.tokens(existing.getTokens())
.ttlSeconds(existing.getTtlSeconds())
.longTermMemoryStrategy(existing.getLongTermMemoryStrategy())
.build();

return putWorkingMemory(sessionId, updated, userId, namespace, null, null);
return putWorkingMemory(sessionId, updated, namespace, userId, null, null);
Copy link

@cursor cursor bot Mar 4, 2026

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Swapped userId and namespace parameters in putWorkingMemory calls

High Severity

All four internal putWorkingMemory calls now pass namespace as the userId parameter and userId as the namespace parameter. The putWorkingMemory signature is (sessionId, memory, userId, namespace, modelName, contextWindowMax), but every changed call site swaps the third and fourth arguments. Since both are @Nullable String, the compiler won't catch this, and API requests will send namespace as user_id and vice versa.

Additional Locations (2)

Fix in Cursor Fix in Web

}

/**
Expand Down Expand Up @@ -434,10 +436,12 @@ public WorkingMemoryResponse addMemoriesToWorkingMemory(
.data(existing.getData())
.context(existing.getContext())
.userId(existing.getUserId())
.tokens(existing.getTokens())
.ttlSeconds(existing.getTtlSeconds())
.longTermMemoryStrategy(existing.getLongTermMemoryStrategy())
.build();

return putWorkingMemory(sessionId, updated, null, namespace, null, null);
return putWorkingMemory(sessionId, updated, namespace, null, null, null);
}

/**
Expand Down Expand Up @@ -505,10 +509,12 @@ public WorkingMemoryResponse updateWorkingMemoryData(
.data(finalData)
.context(existing.getContext())
.userId(existing.getUserId())
.tokens(existing.getTokens())
.ttlSeconds(existing.getTtlSeconds())
.longTermMemoryStrategy(existing.getLongTermMemoryStrategy())
.build();

return putWorkingMemory(sessionId, updated, userId, namespace, null, null);
return putWorkingMemory(sessionId, updated, namespace, userId, null, null);
}

/**
Expand All @@ -521,6 +527,8 @@ public WorkingMemoryResponse updateWorkingMemoryData(
* @param modelName Optional model name for token-based summarization
* @param contextWindowMax Optional context window max tokens
* @param userId Optional user ID
* @param tokens Optional token count for the updated messages (if null, preserves existing token count)
* @param ttlSeconds Optional TTL in seconds to restart the session expiration (if null, preserves existing TTL)
* @return WorkingMemoryResponse with updated memory (potentially summarized if token limit exceeded)
* @throws MemoryClientException if the request fails
*/
Expand All @@ -530,7 +538,9 @@ public WorkingMemoryResponse appendMessagesToWorkingMemory(
@Nullable String namespace,
@Nullable String modelName,
@Nullable Integer contextWindowMax,
@Nullable String userId) throws MemoryClientException {
@Nullable String userId,
@Nullable Integer tokens,
@Nullable Integer ttlSeconds) throws MemoryClientException {
// Get existing memory
WorkingMemoryResult result = getOrCreateWorkingMemory(sessionId, namespace, userId, null, null, null);
WorkingMemoryResponse existing = result.getMemory();
Expand All @@ -541,6 +551,12 @@ public WorkingMemoryResponse appendMessagesToWorkingMemory(
// Append new messages
existingMessages.addAll(messages);

// Determine token count: use provided value, or preserve existing
int tokenCount = tokens != null ? tokens : existing.getTokens();

// Determine TTL: use provided value, or preserve existing
Integer ttl = ttlSeconds != null ? ttlSeconds : existing.getTtlSeconds();

// Create updated working memory
WorkingMemory updated = WorkingMemory.builder()
.sessionId(sessionId)
Expand All @@ -550,10 +566,62 @@ public WorkingMemoryResponse appendMessagesToWorkingMemory(
.data(existing.getData())
.context(existing.getContext())
.userId(userId != null ? userId : existing.getUserId())
.tokens(tokenCount)
.ttlSeconds(ttl)
.longTermMemoryStrategy(existing.getLongTermMemoryStrategy())
.build();

return putWorkingMemory(sessionId, updated, userId, namespace, modelName, contextWindowMax);
return putWorkingMemory(sessionId, updated, namespace, userId, modelName, contextWindowMax);
}

/**
* Append new messages to existing working memory without specifying tokens or TTL.
* Preserves the existing token count and TTL from the session.
*
* @param sessionId The session ID
* @param messages List of messages to append
* @param namespace Optional namespace
* @param modelName Optional model name for token-based summarization
* @param contextWindowMax Optional context window max tokens
* @param userId Optional user ID
* @param tokens Optional token count for the updated messages (if null, preserves existing token count)
* @return WorkingMemoryResponse with updated memory (potentially summarized if token limit exceeded)
* @throws MemoryClientException if the request fails
*/
public WorkingMemoryResponse appendMessagesToWorkingMemory(
@NotNull String sessionId,
@NotNull List<MemoryMessage> messages,
@Nullable String namespace,
@Nullable String modelName,
@Nullable Integer contextWindowMax,
@Nullable String userId,
@Nullable Integer tokens) throws MemoryClientException {
return appendMessagesToWorkingMemory(sessionId, messages, namespace, modelName,
contextWindowMax, userId, tokens, null);
}

/**
* Append new messages to existing working memory without specifying tokens or TTL.
* Preserves the existing token count and TTL from the session.
*
* @param sessionId The session ID
* @param messages List of messages to append
* @param namespace Optional namespace
* @param modelName Optional model name for token-based summarization
* @param contextWindowMax Optional context window max tokens
* @param userId Optional user ID
* @return WorkingMemoryResponse with updated memory (potentially summarized if token limit exceeded)
* @throws MemoryClientException if the request fails
*/
public WorkingMemoryResponse appendMessagesToWorkingMemory(
@NotNull String sessionId,
@NotNull List<MemoryMessage> messages,
@Nullable String namespace,
@Nullable String modelName,
@Nullable Integer contextWindowMax,
@Nullable String userId) throws MemoryClientException {
return appendMessagesToWorkingMemory(sessionId, messages, namespace, modelName,
contextWindowMax, userId, null, null);
}

/**
Expand All @@ -568,7 +636,7 @@ public WorkingMemoryResponse appendMessagesToWorkingMemory(
@NotNull String sessionId,
@NotNull List<MemoryMessage> messages) throws MemoryClientException {
return appendMessagesToWorkingMemory(sessionId, messages, defaultNamespace,
null, null, null);
null, null, null, null, null);
}

// ===== Helper Methods =====
Expand Down
178 changes: 178 additions & 0 deletions ...nt-java/src/test/java/com/redis/agentmemory/integration/WorkingMemoryIntegrationTest.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -298,4 +298,182 @@ void testUpdateWorkingMemoryData() throws Exception {
assertEquals("new_value", response.getData().get("new_field"));
assertEquals("active", response.getData().get("status")); // Should still exist
}

@Test
void testAppendMessagesWithTokens() throws Exception {
String sessionId = "append-tokens-test-" + UUID.randomUUID();
String namespace = "integration-test";

// Create initial working memory with some tokens
WorkingMemory initialMemory = WorkingMemory.builder()
.sessionId(sessionId)
.namespace(namespace)
.messages(Collections.singletonList(
MemoryMessage.builder().role("user").content("Hello").build()))
.tokens(100)
.build();

client.workingMemory().putWorkingMemory(sessionId, initialMemory, namespace, null, null, null);

// Verify initial tokens
WorkingMemoryResponse initial = client.workingMemory()
.getWorkingMemory(sessionId, namespace, null, null, null);
assertEquals(100, initial.getTokens());

// Append messages with updated token count
List<MemoryMessage> newMessages = Collections.singletonList(
MemoryMessage.builder().role("assistant").content("Hi there! How can I help?").build());

WorkingMemoryResponse response = client.workingMemory()
.appendMessagesToWorkingMemory(sessionId, newMessages, namespace, null, null, null, 250);

assertNotNull(response);
assertEquals(250, response.getTokens());
assertTrue(response.getMessages().size() >= 2);

// Verify tokens persisted
WorkingMemoryResponse retrieved = client.workingMemory()
.getWorkingMemory(sessionId, namespace, null, null, null);
assertEquals(250, retrieved.getTokens());
}

@Test
void testAppendMessagesPreservesExistingTokens() throws Exception {
String sessionId = "preserve-tokens-test-" + UUID.randomUUID();
String namespace = "integration-test";

// Create initial working memory with tokens
WorkingMemory initialMemory = WorkingMemory.builder()
.sessionId(sessionId)
.namespace(namespace)
.messages(Collections.singletonList(
MemoryMessage.builder().role("user").content("Hello").build()))
.tokens(150)
.build();

client.workingMemory().putWorkingMemory(sessionId, initialMemory, namespace, null, null, null);

// Append messages WITHOUT specifying tokens (should preserve existing 150)
List<MemoryMessage> newMessages = Collections.singletonList(
MemoryMessage.builder().role("assistant").content("Hi!").build());

WorkingMemoryResponse response = client.workingMemory()
.appendMessagesToWorkingMemory(sessionId, newMessages, namespace, null, null, null);

assertNotNull(response);
assertEquals(150, response.getTokens()); // Should preserve existing
assertTrue(response.getMessages().size() >= 2);

// Verify tokens persisted
WorkingMemoryResponse retrieved = client.workingMemory()
.getWorkingMemory(sessionId, namespace, null, null, null);
assertEquals(150, retrieved.getTokens());
}

@Test
void testAppendMessagesWithTtl() throws Exception {
String sessionId = "append-ttl-test-" + UUID.randomUUID();
String namespace = "integration-test";

// Create initial working memory with TTL
WorkingMemory initialMemory = WorkingMemory.builder()
.sessionId(sessionId)
.namespace(namespace)
.messages(Collections.singletonList(
MemoryMessage.builder().role("user").content("Hello").build()))
.ttlSeconds(1800) // 30 minutes
.build();

client.workingMemory().putWorkingMemory(sessionId, initialMemory, namespace, null, null, null);

// Verify initial TTL
WorkingMemoryResponse initial = client.workingMemory()
.getWorkingMemory(sessionId, namespace, null, null, null);
assertEquals(Integer.valueOf(1800), initial.getTtlSeconds());

// Append messages with new TTL (restart to 1 hour)
List<MemoryMessage> newMessages = Collections.singletonList(
MemoryMessage.builder().role("assistant").content("Hi there!").build());

WorkingMemoryResponse response = client.workingMemory()
.appendMessagesToWorkingMemory(sessionId, newMessages, namespace, null, null, null, null, 3600);

assertNotNull(response);
assertEquals(Integer.valueOf(3600), response.getTtlSeconds());
assertTrue(response.getMessages().size() >= 2);

// Verify TTL persisted
WorkingMemoryResponse retrieved = client.workingMemory()
.getWorkingMemory(sessionId, namespace, null, null, null);
assertEquals(Integer.valueOf(3600), retrieved.getTtlSeconds());
}

@Test
void testAppendMessagesPreservesExistingTtl() throws Exception {
String sessionId = "preserve-ttl-test-" + UUID.randomUUID();
String namespace = "integration-test";

// Create initial working memory with TTL
WorkingMemory initialMemory = WorkingMemory.builder()
.sessionId(sessionId)
.namespace(namespace)
.messages(Collections.singletonList(
MemoryMessage.builder().role("user").content("Hello").build()))
.ttlSeconds(1800) // 30 minutes
.build();

client.workingMemory().putWorkingMemory(sessionId, initialMemory, namespace, null, null, null);

// Append messages WITHOUT specifying TTL (should preserve existing)
List<MemoryMessage> newMessages = Collections.singletonList(
MemoryMessage.builder().role("assistant").content("Hi!").build());

WorkingMemoryResponse response = client.workingMemory()
.appendMessagesToWorkingMemory(sessionId, newMessages, namespace, null, null, null);

assertNotNull(response);
assertEquals(Integer.valueOf(1800), response.getTtlSeconds()); // Should preserve existing
assertTrue(response.getMessages().size() >= 2);

// Verify TTL persisted
WorkingMemoryResponse retrieved = client.workingMemory()
.getWorkingMemory(sessionId, namespace, null, null, null);
assertEquals(Integer.valueOf(1800), retrieved.getTtlSeconds());
}

@Test
void testAppendMessagesWithTokensAndTtl() throws Exception {
String sessionId = "append-tokens-ttl-test-" + UUID.randomUUID();
String namespace = "integration-test";

// Create initial working memory with tokens and TTL
WorkingMemory initialMemory = WorkingMemory.builder()
.sessionId(sessionId)
.namespace(namespace)
.messages(Collections.singletonList(
MemoryMessage.builder().role("user").content("Hello").build()))
.tokens(100)
.ttlSeconds(1800)
.build();

client.workingMemory().putWorkingMemory(sessionId, initialMemory, namespace, null, null, null);

// Append messages with both new tokens and new TTL
List<MemoryMessage> newMessages = Collections.singletonList(
MemoryMessage.builder().role("assistant").content("Hi there!").build());

WorkingMemoryResponse response = client.workingMemory()
.appendMessagesToWorkingMemory(sessionId, newMessages, namespace, null, null, null, 250, 3600);

assertNotNull(response);
assertEquals(250, response.getTokens());
assertEquals(Integer.valueOf(3600), response.getTtlSeconds());
assertTrue(response.getMessages().size() >= 2);

// Verify both persisted
WorkingMemoryResponse retrieved = client.workingMemory()
.getWorkingMemory(sessionId, namespace, null, null, null);
assertEquals(250, retrieved.getTokens());
assertEquals(Integer.valueOf(3600), retrieved.getTtlSeconds());
}
}
Loading
Loading