Skip to content

Commit 345fdd1

Browse files
raphaeldelioraphaeldelio
committed
fix: Preserve tokens and TTL in working memory update methods
1 parent d9d788c commit 345fdd1

File tree

3 files changed

+493
-6
lines changed

3 files changed

+493
-6
lines changed

‎agent-memory-client/agent-memory-client-java/src/main/java/com/redis/agentmemory/services/WorkingMemoryService.java‎

Lines changed: 74 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -365,10 +365,12 @@ public WorkingMemoryResponse setWorkingMemoryData(
365365
.data(data)
366366
.context(existing.getContext())
367367
.userId(existing.getUserId())
368+
.tokens(existing.getTokens())
369+
.ttlSeconds(existing.getTtlSeconds())
368370
.longTermMemoryStrategy(existing.getLongTermMemoryStrategy())
369371
.build();
370372

371-
return putWorkingMemory(sessionId, updated, userId, namespace, null, null);
373+
return putWorkingMemory(sessionId, updated, namespace, userId, null, null);
372374
}
373375

374376
/**
@@ -434,10 +436,12 @@ public WorkingMemoryResponse addMemoriesToWorkingMemory(
434436
.data(existing.getData())
435437
.context(existing.getContext())
436438
.userId(existing.getUserId())
439+
.tokens(existing.getTokens())
440+
.ttlSeconds(existing.getTtlSeconds())
437441
.longTermMemoryStrategy(existing.getLongTermMemoryStrategy())
438442
.build();
439443

440-
return putWorkingMemory(sessionId, updated, null, namespace, null, null);
444+
return putWorkingMemory(sessionId, updated, namespace, null, null, null);
441445
}
442446

443447
/**
@@ -505,10 +509,12 @@ public WorkingMemoryResponse updateWorkingMemoryData(
505509
.data(finalData)
506510
.context(existing.getContext())
507511
.userId(existing.getUserId())
512+
.tokens(existing.getTokens())
513+
.ttlSeconds(existing.getTtlSeconds())
508514
.longTermMemoryStrategy(existing.getLongTermMemoryStrategy())
509515
.build();
510516

511-
return putWorkingMemory(sessionId, updated, userId, namespace, null, null);
517+
return putWorkingMemory(sessionId, updated, namespace, userId, null, null);
512518
}
513519

514520
/**
@@ -521,6 +527,8 @@ public WorkingMemoryResponse updateWorkingMemoryData(
521527
* @param modelName Optional model name for token-based summarization
522528
* @param contextWindowMax Optional context window max tokens
523529
* @param userId Optional user ID
530+
* @param tokens Optional token count for the updated messages (if null, preserves existing token count)
531+
* @param ttlSeconds Optional TTL in seconds to restart the session expiration (if null, preserves existing TTL)
524532
* @return WorkingMemoryResponse with updated memory (potentially summarized if token limit exceeded)
525533
* @throws MemoryClientException if the request fails
526534
*/
@@ -530,7 +538,9 @@ public WorkingMemoryResponse appendMessagesToWorkingMemory(
530538
@Nullable String namespace,
531539
@Nullable String modelName,
532540
@Nullable Integer contextWindowMax,
533-
@Nullable String userId) throws MemoryClientException {
541+
@Nullable String userId,
542+
@Nullable Integer tokens,
543+
@Nullable Integer ttlSeconds) throws MemoryClientException {
534544
// Get existing memory
535545
WorkingMemoryResult result = getOrCreateWorkingMemory(sessionId, namespace, userId, null, null, null);
536546
WorkingMemoryResponse existing = result.getMemory();
@@ -541,6 +551,12 @@ public WorkingMemoryResponse appendMessagesToWorkingMemory(
541551
// Append new messages
542552
existingMessages.addAll(messages);
543553

554+
// Determine token count: use provided value, or preserve existing
555+
int tokenCount = tokens != null ? tokens : existing.getTokens();
556+
557+
// Determine TTL: use provided value, or preserve existing
558+
Integer ttl = ttlSeconds != null ? ttlSeconds : existing.getTtlSeconds();
559+
544560
// Create updated working memory
545561
WorkingMemory updated = WorkingMemory.builder()
546562
.sessionId(sessionId)
@@ -550,10 +566,62 @@ public WorkingMemoryResponse appendMessagesToWorkingMemory(
550566
.data(existing.getData())
551567
.context(existing.getContext())
552568
.userId(userId != null ? userId : existing.getUserId())
569+
.tokens(tokenCount)
570+
.ttlSeconds(ttl)
553571
.longTermMemoryStrategy(existing.getLongTermMemoryStrategy())
554572
.build();
555573

556-
return putWorkingMemory(sessionId, updated, userId, namespace, modelName, contextWindowMax);
574+
return putWorkingMemory(sessionId, updated, namespace, userId, modelName, contextWindowMax);
575+
}
576+
577+
/**
578+
* Append new messages to existing working memory without specifying tokens or TTL.
579+
* Preserves the existing token count and TTL from the session.
580+
*
581+
* @param sessionId The session ID
582+
* @param messages List of messages to append
583+
* @param namespace Optional namespace
584+
* @param modelName Optional model name for token-based summarization
585+
* @param contextWindowMax Optional context window max tokens
586+
* @param userId Optional user ID
587+
* @param tokens Optional token count for the updated messages (if null, preserves existing token count)
588+
* @return WorkingMemoryResponse with updated memory (potentially summarized if token limit exceeded)
589+
* @throws MemoryClientException if the request fails
590+
*/
591+
public WorkingMemoryResponse appendMessagesToWorkingMemory(
592+
@NotNull String sessionId,
593+
@NotNull List<MemoryMessage> messages,
594+
@Nullable String namespace,
595+
@Nullable String modelName,
596+
@Nullable Integer contextWindowMax,
597+
@Nullable String userId,
598+
@Nullable Integer tokens) throws MemoryClientException {
599+
return appendMessagesToWorkingMemory(sessionId, messages, namespace, modelName,
600+
contextWindowMax, userId, tokens, null);
601+
}
602+
603+
/**
604+
* Append new messages to existing working memory without specifying tokens or TTL.
605+
* Preserves the existing token count and TTL from the session.
606+
*
607+
* @param sessionId The session ID
608+
* @param messages List of messages to append
609+
* @param namespace Optional namespace
610+
* @param modelName Optional model name for token-based summarization
611+
* @param contextWindowMax Optional context window max tokens
612+
* @param userId Optional user ID
613+
* @return WorkingMemoryResponse with updated memory (potentially summarized if token limit exceeded)
614+
* @throws MemoryClientException if the request fails
615+
*/
616+
public WorkingMemoryResponse appendMessagesToWorkingMemory(
617+
@NotNull String sessionId,
618+
@NotNull List<MemoryMessage> messages,
619+
@Nullable String namespace,
620+
@Nullable String modelName,
621+
@Nullable Integer contextWindowMax,
622+
@Nullable String userId) throws MemoryClientException {
623+
return appendMessagesToWorkingMemory(sessionId, messages, namespace, modelName,
624+
contextWindowMax, userId, null, null);
557625
}
558626

559627
/**
@@ -568,7 +636,7 @@ public WorkingMemoryResponse appendMessagesToWorkingMemory(
568636
@NotNull String sessionId,
569637
@NotNull List<MemoryMessage> messages) throws MemoryClientException {
570638
return appendMessagesToWorkingMemory(sessionId, messages, defaultNamespace,
571-
null, null, null);
639+
null, null, null, null, null);
572640
}
573641

574642
// ===== Helper Methods =====

‎agent-memory-client/agent-memory-client-java/src/test/java/com/redis/agentmemory/integration/WorkingMemoryIntegrationTest.java‎

Lines changed: 178 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -298,4 +298,182 @@ void testUpdateWorkingMemoryData() throws Exception {
298298
assertEquals("new_value", response.getData().get("new_field"));
299299
assertEquals("active", response.getData().get("status")); // Should still exist
300300
}
301+
302+
@Test
303+
void testAppendMessagesWithTokens() throws Exception {
304+
String sessionId = "append-tokens-test-" + UUID.randomUUID();
305+
String namespace = "integration-test";
306+
307+
// Create initial working memory with some tokens
308+
WorkingMemory initialMemory = WorkingMemory.builder()
309+
.sessionId(sessionId)
310+
.namespace(namespace)
311+
.messages(Collections.singletonList(
312+
MemoryMessage.builder().role("user").content("Hello").build()))
313+
.tokens(100)
314+
.build();
315+
316+
client.workingMemory().putWorkingMemory(sessionId, initialMemory, namespace, null, null, null);
317+
318+
// Verify initial tokens
319+
WorkingMemoryResponse initial = client.workingMemory()
320+
.getWorkingMemory(sessionId, namespace, null, null, null);
321+
assertEquals(100, initial.getTokens());
322+
323+
// Append messages with updated token count
324+
List<MemoryMessage> newMessages = Collections.singletonList(
325+
MemoryMessage.builder().role("assistant").content("Hi there! How can I help?").build());
326+
327+
WorkingMemoryResponse response = client.workingMemory()
328+
.appendMessagesToWorkingMemory(sessionId, newMessages, namespace, null, null, null, 250);
329+
330+
assertNotNull(response);
331+
assertEquals(250, response.getTokens());
332+
assertTrue(response.getMessages().size() >= 2);
333+
334+
// Verify tokens persisted
335+
WorkingMemoryResponse retrieved = client.workingMemory()
336+
.getWorkingMemory(sessionId, namespace, null, null, null);
337+
assertEquals(250, retrieved.getTokens());
338+
}
339+
340+
@Test
341+
void testAppendMessagesPreservesExistingTokens() throws Exception {
342+
String sessionId = "preserve-tokens-test-" + UUID.randomUUID();
343+
String namespace = "integration-test";
344+
345+
// Create initial working memory with tokens
346+
WorkingMemory initialMemory = WorkingMemory.builder()
347+
.sessionId(sessionId)
348+
.namespace(namespace)
349+
.messages(Collections.singletonList(
350+
MemoryMessage.builder().role("user").content("Hello").build()))
351+
.tokens(150)
352+
.build();
353+
354+
client.workingMemory().putWorkingMemory(sessionId, initialMemory, namespace, null, null, null);
355+
356+
// Append messages WITHOUT specifying tokens (should preserve existing 150)
357+
List<MemoryMessage> newMessages = Collections.singletonList(
358+
MemoryMessage.builder().role("assistant").content("Hi!").build());
359+
360+
WorkingMemoryResponse response = client.workingMemory()
361+
.appendMessagesToWorkingMemory(sessionId, newMessages, namespace, null, null, null);
362+
363+
assertNotNull(response);
364+
assertEquals(150, response.getTokens()); // Should preserve existing
365+
assertTrue(response.getMessages().size() >= 2);
366+
367+
// Verify tokens persisted
368+
WorkingMemoryResponse retrieved = client.workingMemory()
369+
.getWorkingMemory(sessionId, namespace, null, null, null);
370+
assertEquals(150, retrieved.getTokens());
371+
}
372+
373+
@Test
374+
void testAppendMessagesWithTtl() throws Exception {
375+
String sessionId = "append-ttl-test-" + UUID.randomUUID();
376+
String namespace = "integration-test";
377+
378+
// Create initial working memory with TTL
379+
WorkingMemory initialMemory = WorkingMemory.builder()
380+
.sessionId(sessionId)
381+
.namespace(namespace)
382+
.messages(Collections.singletonList(
383+
MemoryMessage.builder().role("user").content("Hello").build()))
384+
.ttlSeconds(1800) // 30 minutes
385+
.build();
386+
387+
client.workingMemory().putWorkingMemory(sessionId, initialMemory, namespace, null, null, null);
388+
389+
// Verify initial TTL
390+
WorkingMemoryResponse initial = client.workingMemory()
391+
.getWorkingMemory(sessionId, namespace, null, null, null);
392+
assertEquals(Integer.valueOf(1800), initial.getTtlSeconds());
393+
394+
// Append messages with new TTL (restart to 1 hour)
395+
List<MemoryMessage> newMessages = Collections.singletonList(
396+
MemoryMessage.builder().role("assistant").content("Hi there!").build());
397+
398+
WorkingMemoryResponse response = client.workingMemory()
399+
.appendMessagesToWorkingMemory(sessionId, newMessages, namespace, null, null, null, null, 3600);
400+
401+
assertNotNull(response);
402+
assertEquals(Integer.valueOf(3600), response.getTtlSeconds());
403+
assertTrue(response.getMessages().size() >= 2);
404+
405+
// Verify TTL persisted
406+
WorkingMemoryResponse retrieved = client.workingMemory()
407+
.getWorkingMemory(sessionId, namespace, null, null, null);
408+
assertEquals(Integer.valueOf(3600), retrieved.getTtlSeconds());
409+
}
410+
411+
@Test
412+
void testAppendMessagesPreservesExistingTtl() throws Exception {
413+
String sessionId = "preserve-ttl-test-" + UUID.randomUUID();
414+
String namespace = "integration-test";
415+
416+
// Create initial working memory with TTL
417+
WorkingMemory initialMemory = WorkingMemory.builder()
418+
.sessionId(sessionId)
419+
.namespace(namespace)
420+
.messages(Collections.singletonList(
421+
MemoryMessage.builder().role("user").content("Hello").build()))
422+
.ttlSeconds(1800) // 30 minutes
423+
.build();
424+
425+
client.workingMemory().putWorkingMemory(sessionId, initialMemory, namespace, null, null, null);
426+
427+
// Append messages WITHOUT specifying TTL (should preserve existing)
428+
List<MemoryMessage> newMessages = Collections.singletonList(
429+
MemoryMessage.builder().role("assistant").content("Hi!").build());
430+
431+
WorkingMemoryResponse response = client.workingMemory()
432+
.appendMessagesToWorkingMemory(sessionId, newMessages, namespace, null, null, null);
433+
434+
assertNotNull(response);
435+
assertEquals(Integer.valueOf(1800), response.getTtlSeconds()); // Should preserve existing
436+
assertTrue(response.getMessages().size() >= 2);
437+
438+
// Verify TTL persisted
439+
WorkingMemoryResponse retrieved = client.workingMemory()
440+
.getWorkingMemory(sessionId, namespace, null, null, null);
441+
assertEquals(Integer.valueOf(1800), retrieved.getTtlSeconds());
442+
}
443+
444+
@Test
445+
void testAppendMessagesWithTokensAndTtl() throws Exception {
446+
String sessionId = "append-tokens-ttl-test-" + UUID.randomUUID();
447+
String namespace = "integration-test";
448+
449+
// Create initial working memory with tokens and TTL
450+
WorkingMemory initialMemory = WorkingMemory.builder()
451+
.sessionId(sessionId)
452+
.namespace(namespace)
453+
.messages(Collections.singletonList(
454+
MemoryMessage.builder().role("user").content("Hello").build()))
455+
.tokens(100)
456+
.ttlSeconds(1800)
457+
.build();
458+
459+
client.workingMemory().putWorkingMemory(sessionId, initialMemory, namespace, null, null, null);
460+
461+
// Append messages with both new tokens and new TTL
462+
List<MemoryMessage> newMessages = Collections.singletonList(
463+
MemoryMessage.builder().role("assistant").content("Hi there!").build());
464+
465+
WorkingMemoryResponse response = client.workingMemory()
466+
.appendMessagesToWorkingMemory(sessionId, newMessages, namespace, null, null, null, 250, 3600);
467+
468+
assertNotNull(response);
469+
assertEquals(250, response.getTokens());
470+
assertEquals(Integer.valueOf(3600), response.getTtlSeconds());
471+
assertTrue(response.getMessages().size() >= 2);
472+
473+
// Verify both persisted
474+
WorkingMemoryResponse retrieved = client.workingMemory()
475+
.getWorkingMemory(sessionId, namespace, null, null, null);
476+
assertEquals(250, retrieved.getTokens());
477+
assertEquals(Integer.valueOf(3600), retrieved.getTtlSeconds());
478+
}
301479
}

0 commit comments

Comments
 (0)