tkmind/tkmind_go
1
0
Fork 0
Code Issues Pull Requests Actions 4 Packages Projects Releases Wiki Activity

Nest TODO State in session data (#4361)

Browse Source 
Co-authored-by: Alex Hancock <alexhancock@block.xyz>
This commit is contained in:
David Katz
2025-08-28 14:11:29 -04:00
committed by GitHub
parent 7879445c4b
commit ef57f5062d
11 changed files with 294 additions and 45 deletions
Download Patch File Download Diff File Expand all files Collapse all files
crates/goose/tests/test_support.rs
+1 -1
View File
@@ -411,6 +411,6 @@ pub fn create_test_session_metadata(message_count: usize, working_dir: &str) ->
accumulated_total_tokens: Some(100),
accumulated_input_tokens: Some(50),
accumulated_output_tokens: Some(50),
todo_content: None,
extension_data: Default::default(),
}
}
crates/goose/tests/todo_session_integration.rs
+77 -24
View File
@@ -160,7 +160,10 @@ async fn test_todo_add_persists_to_session() {
// Since we're using a mock provider, we can't test the actual TODO content
// but we can verify the metadata structure is correct
assert!(metadata.todo_content.is_some() || metadata.todo_content.is_none());
assert!(
metadata.extension_data.extension_states.is_empty()
|| !metadata.extension_data.extension_states.is_empty()
);
}
#[tokio::test]
@@ -172,7 +175,11 @@ async fn test_todo_list_reads_from_session() {
// Pre-populate session with TODO content
let session_path = goose::session::storage::get_path(session_id.clone()).unwrap();
let mut metadata = SessionMetadata::default();
metadata.todo_content = Some("- Task 1\n- Task 2\n- Task 3".to_string());
use goose::session::extension_data::{ExtensionState, TodoState};
let todo_state = TodoState::new("- Task 1\n- Task 2\n- Task 3".to_string());
todo_state
.to_extension_data(&mut metadata.extension_data)
.unwrap();
goose::session::storage::update_metadata(&session_path, &metadata)
.await
.unwrap();
@@ -206,21 +213,27 @@ async fn test_todo_list_reads_from_session() {
// Verify the TODO content is still in session
let metadata_after = goose::session::storage::read_metadata(&session_path).unwrap();
let todo_state_after = TodoState::from_extension_data(&metadata_after.extension_data);
assert!(todo_state_after.is_some());
assert_eq!(
metadata_after.todo_content,
Some("- Task 1\n- Task 2\n- Task 3".to_string())
todo_state_after.unwrap().content,
"- Task 1\n- Task 2\n- Task 3".to_string()
);
}
#[tokio::test]
async fn test_todo_isolation_between_sessions() {
use goose::session::extension_data::{ExtensionState, TodoState};
let session1_id = session::Identifier::Name(format!("test_session_{}", Uuid::new_v4()));
let session2_id = session::Identifier::Name(format!("test_session_{}", Uuid::new_v4()));
// Add TODO to session1
let session1_path = goose::session::storage::get_path(session1_id.clone()).unwrap();
let mut metadata1 = SessionMetadata::default();
metadata1.todo_content = Some("Session 1 tasks".to_string());
let todo_state1 = TodoState::new("Session 1 tasks".to_string());
todo_state1
.to_extension_data(&mut metadata1.extension_data)
.unwrap();
goose::session::storage::update_metadata(&session1_path, &metadata1)
.await
.unwrap();
@@ -228,7 +241,10 @@ async fn test_todo_isolation_between_sessions() {
// Add different TODO to session2
let session2_path = goose::session::storage::get_path(session2_id.clone()).unwrap();
let mut metadata2 = SessionMetadata::default();
metadata2.todo_content = Some("Session 2 tasks".to_string());
let todo_state2 = TodoState::new("Session 2 tasks".to_string());
todo_state2
.to_extension_data(&mut metadata2.extension_data)
.unwrap();
goose::session::storage::update_metadata(&session2_path, &metadata2)
.await
.unwrap();
@@ -237,12 +253,16 @@ async fn test_todo_isolation_between_sessions() {
let metadata1_read = goose::session::storage::read_metadata(&session1_path).unwrap();
let metadata2_read = goose::session::storage::read_metadata(&session2_path).unwrap();
assert_eq!(metadata1_read.todo_content.unwrap(), "Session 1 tasks");
assert_eq!(metadata2_read.todo_content.unwrap(), "Session 2 tasks");
let todo1 = TodoState::from_extension_data(&metadata1_read.extension_data).unwrap();
let todo2 = TodoState::from_extension_data(&metadata2_read.extension_data).unwrap();
assert_eq!(todo1.content, "Session 1 tasks");
assert_eq!(todo2.content, "Session 2 tasks");
}
#[tokio::test]
async fn test_todo_clear_removes_from_session() {
use goose::session::extension_data::{ExtensionState, TodoState};
let temp_dir = create_test_session_dir().await;
let session_id = session::Identifier::Name(format!("test_session_{}", Uuid::new_v4()));
let agent = create_test_agent_with_mock_provider().await;
@@ -250,7 +270,10 @@ async fn test_todo_clear_removes_from_session() {
// Pre-populate session with TODO content
let session_path = goose::session::storage::get_path(session_id.clone()).unwrap();
let mut metadata = SessionMetadata::default();
metadata.todo_content = Some("- Task to clear".to_string());
let todo_state = TodoState::new("- Task to clear".to_string());
todo_state
.to_extension_data(&mut metadata.extension_data)
.unwrap();
goose::session::storage::update_metadata(&session_path, &metadata)
.await
.unwrap();
@@ -280,18 +303,23 @@ async fn test_todo_clear_removes_from_session() {
// With mock provider, the TODO won't actually be cleared via tool calls
// but we can verify the structure is correct
let metadata_after = goose::session::storage::read_metadata(&session_path).unwrap();
assert!(metadata_after.todo_content.is_some()); // Will still have the original content with mock
let todo_state_after = TodoState::from_extension_data(&metadata_after.extension_data);
assert!(todo_state_after.is_some()); // Will still have the original content with mock
}
#[tokio::test]
async fn test_todo_persistence_across_agent_instances() {
use goose::session::extension_data::{ExtensionState, TodoState};
let session_id = session::Identifier::Name(format!("test_session_{}", Uuid::new_v4()));
// First agent instance adds TODO
{
let session_path = goose::session::storage::get_path(session_id.clone()).unwrap();
let mut metadata = SessionMetadata::default();
metadata.todo_content = Some("Persistent task".to_string());
let todo_state = TodoState::new("Persistent task".to_string());
todo_state
.to_extension_data(&mut metadata.extension_data)
.unwrap();
goose::session::storage::update_metadata(&session_path, &metadata)
.await
.unwrap();
@@ -301,13 +329,14 @@ async fn test_todo_persistence_across_agent_instances() {
{
let session_path = goose::session::storage::get_path(session_id.clone()).unwrap();
let metadata = goose::session::storage::read_metadata(&session_path).unwrap();
assert_eq!(metadata.todo_content.unwrap(), "Persistent task");
let todo_state = TodoState::from_extension_data(&metadata.extension_data).unwrap();
assert_eq!(todo_state.content, "Persistent task");
}
}
#[tokio::test]
async fn test_todo_max_chars_limit() {
use goose::session::extension_data::{ExtensionState, TodoState};
let session_id = session::Identifier::Name(format!("test_session_{}", Uuid::new_v4()));
// Set a small limit for testing
@@ -318,7 +347,10 @@ async fn test_todo_max_chars_limit() {
// Try to set content that exceeds the limit
let long_content = "x".repeat(100);
metadata.todo_content = Some(long_content.clone());
let todo_state = TodoState::new(long_content.clone());
todo_state
.to_extension_data(&mut metadata.extension_data)
.unwrap();
// This should succeed at the storage level (storage doesn't enforce limits)
goose::session::storage::update_metadata(&session_path, &metadata)
@@ -334,6 +366,7 @@ async fn test_todo_max_chars_limit() {
#[tokio::test]
async fn test_todo_with_special_characters() {
use goose::session::extension_data::{ExtensionState, TodoState};
let session_id = session::Identifier::Name(format!("test_session_{}", Uuid::new_v4()));
let session_path = goose::session::storage::get_path(session_id.clone()).unwrap();
@@ -350,18 +383,23 @@ async fn test_todo_with_special_characters() {
- Task with tab separation
"#;
metadata.todo_content = Some(special_content.to_string());
let todo_state = TodoState::new(special_content.to_string());
todo_state
.to_extension_data(&mut metadata.extension_data)
.unwrap();
goose::session::storage::update_metadata(&session_path, &metadata)
.await
.unwrap();
// Read back and verify
let metadata_read = goose::session::storage::read_metadata(&session_path).unwrap();
assert_eq!(metadata_read.todo_content.unwrap(), special_content);
let todo_state_read = TodoState::from_extension_data(&metadata_read.extension_data).unwrap();
assert_eq!(todo_state_read.content, special_content);
}
#[tokio::test]
async fn test_todo_concurrent_access() {
use goose::session::extension_data::{ExtensionState, TodoState};
let session_id = session::Identifier::Name(format!("test_session_{}", Uuid::new_v4()));
// Spawn multiple concurrent TODO operations
@@ -375,8 +413,13 @@ async fn test_todo_concurrent_access() {
let mut metadata = goose::session::storage::read_metadata(&session_path)
.unwrap_or_else(|_| SessionMetadata::default());
let current_content = metadata.todo_content.unwrap_or_default();
metadata.todo_content = Some(format!("{}\n- Task {}", current_content, i));
let current_content = TodoState::from_extension_data(&metadata.extension_data)
.map(|t| t.content)
.unwrap_or_default();
let new_todo = TodoState::new(format!("{}\n- Task {}", current_content, i));
new_todo
.to_extension_data(&mut metadata.extension_data)
.unwrap();
goose::session::storage::update_metadata(&session_path, &metadata).await
});
@@ -392,25 +435,28 @@ async fn test_todo_concurrent_access() {
// Verify final state contains at least one task
let session_path = goose::session::storage::get_path(session_id).unwrap();
let metadata = goose::session::storage::read_metadata(&session_path).unwrap();
let todo_content = metadata.todo_content.unwrap();
let todo_state = TodoState::from_extension_data(&metadata.extension_data).unwrap();
// Should contain at least one task (concurrent writes may overwrite)
assert!(todo_content.contains("Task"));
assert!(todo_state.content.contains("Task"));
}
#[tokio::test]
async fn test_todo_empty_session_returns_empty() {
use goose::session::extension_data::{ExtensionState, TodoState};
let session_id = session::Identifier::Name(format!("test_session_{}", Uuid::new_v4()));
let session_path = goose::session::storage::get_path(session_id.clone()).unwrap();
let metadata = goose::session::storage::read_metadata(&session_path)
.unwrap_or_else(|_| SessionMetadata::default());
assert!(metadata.todo_content.is_none() || metadata.todo_content.as_ref().unwrap().is_empty());
let todo_state = TodoState::from_extension_data(&metadata.extension_data);
assert!(todo_state.is_none() || todo_state.unwrap().content.is_empty());
}
#[tokio::test]
async fn test_todo_update_preserves_other_metadata() {
use goose::session::extension_data::{ExtensionState, TodoState};
let session_id = session::Identifier::Name(format!("test_session_{}", Uuid::new_v4()));
let session_path = goose::session::storage::get_path(session_id.clone()).unwrap();
@@ -420,14 +466,20 @@ async fn test_todo_update_preserves_other_metadata() {
metadata.message_count = 5;
metadata.description = "Test session".to_string();
metadata.total_tokens = Some(1000);
metadata.todo_content = Some("Initial TODO".to_string());
let todo_state = TodoState::new("Initial TODO".to_string());
todo_state
.to_extension_data(&mut metadata.extension_data)
.unwrap();
goose::session::storage::update_metadata(&session_path, &metadata)
.await
.unwrap();
// Update only TODO content
metadata.todo_content = Some("Updated TODO".to_string());
let todo_state_updated = TodoState::new("Updated TODO".to_string());
todo_state_updated
.to_extension_data(&mut metadata.extension_data)
.unwrap();
goose::session::storage::update_metadata(&session_path, &metadata)
.await
.unwrap();
@@ -437,5 +489,6 @@ async fn test_todo_update_preserves_other_metadata() {
assert_eq!(metadata_read.message_count, 5);
assert_eq!(metadata_read.description, "Test session");
assert_eq!(metadata_read.total_tokens, Some(1000));
assert_eq!(metadata_read.todo_content, Some("Updated TODO".to_string()));
let todo_state_read = TodoState::from_extension_data(&metadata_read.extension_data).unwrap();
assert_eq!(todo_state_read.content, "Updated TODO");
}