Nest TODO State in session data (#4361)
Co-authored-by: Alex Hancock <alexhancock@block.xyz>
This commit is contained in:
@@ -445,6 +445,7 @@ impl<'__s> ToSchema<'__s> for AnnotatedSchema {
|
|||||||
ModelInfo,
|
ModelInfo,
|
||||||
SessionInfo,
|
SessionInfo,
|
||||||
SessionMetadata,
|
SessionMetadata,
|
||||||
|
goose::session::ExtensionData,
|
||||||
super::routes::schedule::CreateScheduleRequest,
|
super::routes::schedule::CreateScheduleRequest,
|
||||||
super::routes::schedule::UpdateScheduleRequest,
|
super::routes::schedule::UpdateScheduleRequest,
|
||||||
super::routes::schedule::KillJobResponse,
|
super::routes::schedule::KillJobResponse,
|
||||||
|
|||||||
@@ -41,6 +41,7 @@ use crate::providers::errors::ProviderError;
|
|||||||
use crate::recipe::{Author, Recipe, Response, Settings, SubRecipe};
|
use crate::recipe::{Author, Recipe, Response, Settings, SubRecipe};
|
||||||
use crate::scheduler_trait::SchedulerTrait;
|
use crate::scheduler_trait::SchedulerTrait;
|
||||||
use crate::session;
|
use crate::session;
|
||||||
|
use crate::session::extension_data::ExtensionState;
|
||||||
use crate::tool_monitor::{ToolCall, ToolMonitor};
|
use crate::tool_monitor::{ToolCall, ToolMonitor};
|
||||||
use crate::utils::is_token_cancelled;
|
use crate::utils::is_token_cancelled;
|
||||||
use mcp_core::ToolResult;
|
use mcp_core::ToolResult;
|
||||||
@@ -494,7 +495,10 @@ impl Agent {
|
|||||||
let todo_content = if let Some(path) = session_file_path {
|
let todo_content = if let Some(path) = session_file_path {
|
||||||
session::storage::read_metadata(&path)
|
session::storage::read_metadata(&path)
|
||||||
.ok()
|
.ok()
|
||||||
.and_then(|m| m.todo_content)
|
.and_then(|m| {
|
||||||
|
session::TodoState::from_extension_data(&m.extension_data)
|
||||||
|
.map(|state| state.content)
|
||||||
|
})
|
||||||
.unwrap_or_default()
|
.unwrap_or_default()
|
||||||
} else {
|
} else {
|
||||||
String::new()
|
String::new()
|
||||||
@@ -531,7 +535,11 @@ impl Agent {
|
|||||||
match session::storage::get_path(session_config.id.clone()) {
|
match session::storage::get_path(session_config.id.clone()) {
|
||||||
Ok(path) => match session::storage::read_metadata(&path) {
|
Ok(path) => match session::storage::read_metadata(&path) {
|
||||||
Ok(mut metadata) => {
|
Ok(mut metadata) => {
|
||||||
metadata.todo_content = Some(content);
|
let todo_state = session::TodoState::new(content);
|
||||||
|
todo_state
|
||||||
|
.to_extension_data(&mut metadata.extension_data)
|
||||||
|
.ok();
|
||||||
|
|
||||||
let path_clone = path.clone();
|
let path_clone = path.clone();
|
||||||
let metadata_clone = metadata.clone();
|
let metadata_clone = metadata.clone();
|
||||||
let update_result = tokio::task::spawn(async move {
|
let update_result = tokio::task::spawn(async move {
|
||||||
|
|||||||
@@ -269,7 +269,7 @@ mod tests {
|
|||||||
accumulated_total_tokens: Some(100),
|
accumulated_total_tokens: Some(100),
|
||||||
accumulated_input_tokens: Some(50),
|
accumulated_input_tokens: Some(50),
|
||||||
accumulated_output_tokens: Some(50),
|
accumulated_output_tokens: Some(50),
|
||||||
todo_content: None,
|
extension_data: crate::session::ExtensionData::new(),
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
|
|||||||
@@ -1298,7 +1298,7 @@ async fn run_scheduled_job_internal(
|
|||||||
accumulated_total_tokens: None,
|
accumulated_total_tokens: None,
|
||||||
accumulated_input_tokens: None,
|
accumulated_input_tokens: None,
|
||||||
accumulated_output_tokens: None,
|
accumulated_output_tokens: None,
|
||||||
todo_content: None,
|
extension_data: crate::session::ExtensionData::new(),
|
||||||
};
|
};
|
||||||
if let Err(e_fb) = crate::session::storage::save_messages_with_metadata(
|
if let Err(e_fb) = crate::session::storage::save_messages_with_metadata(
|
||||||
&session_file_path,
|
&session_file_path,
|
||||||
|
|||||||
@@ -0,0 +1,173 @@
|
|||||||
|
// Extension data management for sessions
|
||||||
|
// Provides a simple way to store extension-specific data with versioned keys
|
||||||
|
|
||||||
|
use anyhow::Result;
|
||||||
|
use serde::{Deserialize, Serialize};
|
||||||
|
use serde_json::Value;
|
||||||
|
use std::collections::HashMap;
|
||||||
|
use utoipa::ToSchema;
|
||||||
|
|
||||||
|
/// Extension data containing all extension states
|
||||||
|
/// Keys are in format "extension_name.version" (e.g., "todo.v0")
|
||||||
|
#[derive(Debug, Clone, Serialize, Deserialize, Default, ToSchema)]
|
||||||
|
pub struct ExtensionData {
|
||||||
|
#[serde(flatten)]
|
||||||
|
pub extension_states: HashMap<String, Value>,
|
||||||
|
}
|
||||||
|
|
||||||
|
impl ExtensionData {
|
||||||
|
/// Create a new empty ExtensionData
|
||||||
|
pub fn new() -> Self {
|
||||||
|
Self {
|
||||||
|
extension_states: HashMap::new(),
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Get extension state for a specific extension and version
|
||||||
|
pub fn get_extension_state(&self, extension_name: &str, version: &str) -> Option<&Value> {
|
||||||
|
let key = format!("{}.{}", extension_name, version);
|
||||||
|
self.extension_states.get(&key)
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Set extension state for a specific extension and version
|
||||||
|
pub fn set_extension_state(&mut self, extension_name: &str, version: &str, state: Value) {
|
||||||
|
let key = format!("{}.{}", extension_name, version);
|
||||||
|
self.extension_states.insert(key, state);
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Helper trait for extension-specific state management
|
||||||
|
pub trait ExtensionState: Sized + Serialize + for<'de> Deserialize<'de> {
|
||||||
|
/// The name of the extension
|
||||||
|
const EXTENSION_NAME: &'static str;
|
||||||
|
|
||||||
|
/// The version of the extension state format
|
||||||
|
const VERSION: &'static str;
|
||||||
|
|
||||||
|
/// Convert from JSON value
|
||||||
|
fn from_value(value: &Value) -> Result<Self> {
|
||||||
|
serde_json::from_value(value.clone()).map_err(|e| {
|
||||||
|
anyhow::anyhow!(
|
||||||
|
"Failed to deserialize {} state: {}",
|
||||||
|
Self::EXTENSION_NAME,
|
||||||
|
e
|
||||||
|
)
|
||||||
|
})
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Convert to JSON value
|
||||||
|
fn to_value(&self) -> Result<Value> {
|
||||||
|
serde_json::to_value(self).map_err(|e| {
|
||||||
|
anyhow::anyhow!("Failed to serialize {} state: {}", Self::EXTENSION_NAME, e)
|
||||||
|
})
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Get state from extension data
|
||||||
|
fn from_extension_data(extension_data: &ExtensionData) -> Option<Self> {
|
||||||
|
extension_data
|
||||||
|
.get_extension_state(Self::EXTENSION_NAME, Self::VERSION)
|
||||||
|
.and_then(|v| Self::from_value(v).ok())
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Save state to extension data
|
||||||
|
fn to_extension_data(&self, extension_data: &mut ExtensionData) -> Result<()> {
|
||||||
|
let value = self.to_value()?;
|
||||||
|
extension_data.set_extension_state(Self::EXTENSION_NAME, Self::VERSION, value);
|
||||||
|
Ok(())
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
/// TODO extension state implementation
|
||||||
|
#[derive(Debug, Clone, Serialize, Deserialize)]
|
||||||
|
pub struct TodoState {
|
||||||
|
pub content: String,
|
||||||
|
}
|
||||||
|
|
||||||
|
impl ExtensionState for TodoState {
|
||||||
|
const EXTENSION_NAME: &'static str = "todo";
|
||||||
|
const VERSION: &'static str = "v0";
|
||||||
|
}
|
||||||
|
|
||||||
|
impl TodoState {
|
||||||
|
/// Create a new TODO state
|
||||||
|
pub fn new(content: String) -> Self {
|
||||||
|
Self { content }
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
#[cfg(test)]
|
||||||
|
mod tests {
|
||||||
|
use super::*;
|
||||||
|
use serde_json::json;
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
fn test_extension_data_basic_operations() {
|
||||||
|
let mut extension_data = ExtensionData::new();
|
||||||
|
|
||||||
|
// Test setting and getting extension state
|
||||||
|
let todo_state = json!({"content": "- Task 1\n- Task 2"});
|
||||||
|
extension_data.set_extension_state("todo", "v0", todo_state.clone());
|
||||||
|
|
||||||
|
assert_eq!(
|
||||||
|
extension_data.get_extension_state("todo", "v0"),
|
||||||
|
Some(&todo_state)
|
||||||
|
);
|
||||||
|
assert_eq!(extension_data.get_extension_state("todo", "v1"), None);
|
||||||
|
}
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
fn test_multiple_extension_states() {
|
||||||
|
let mut extension_data = ExtensionData::new();
|
||||||
|
|
||||||
|
// Add multiple extension states
|
||||||
|
extension_data.set_extension_state("todo", "v0", json!("TODO content"));
|
||||||
|
extension_data.set_extension_state("memory", "v1", json!({"items": ["item1", "item2"]}));
|
||||||
|
extension_data.set_extension_state("config", "v2", json!({"setting": true}));
|
||||||
|
|
||||||
|
// Check all states exist
|
||||||
|
assert_eq!(extension_data.extension_states.len(), 3);
|
||||||
|
assert!(extension_data.get_extension_state("todo", "v0").is_some());
|
||||||
|
assert!(extension_data.get_extension_state("memory", "v1").is_some());
|
||||||
|
assert!(extension_data.get_extension_state("config", "v2").is_some());
|
||||||
|
}
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
fn test_todo_state_trait() {
|
||||||
|
let mut extension_data = ExtensionData::new();
|
||||||
|
|
||||||
|
// Create and save TODO state
|
||||||
|
let todo = TodoState::new("- Task 1\n- Task 2".to_string());
|
||||||
|
todo.to_extension_data(&mut extension_data).unwrap();
|
||||||
|
|
||||||
|
// Retrieve TODO state
|
||||||
|
let retrieved = TodoState::from_extension_data(&extension_data);
|
||||||
|
assert!(retrieved.is_some());
|
||||||
|
assert_eq!(retrieved.unwrap().content, "- Task 1\n- Task 2");
|
||||||
|
}
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
fn test_extension_data_serialization() {
|
||||||
|
let mut extension_data = ExtensionData::new();
|
||||||
|
extension_data.set_extension_state("todo", "v0", json!("TODO content"));
|
||||||
|
extension_data.set_extension_state("memory", "v1", json!({"key": "value"}));
|
||||||
|
|
||||||
|
// Serialize to JSON
|
||||||
|
let json = serde_json::to_value(&extension_data).unwrap();
|
||||||
|
|
||||||
|
// Check the structure
|
||||||
|
assert!(json.is_object());
|
||||||
|
assert_eq!(json.get("todo.v0"), Some(&json!("TODO content")));
|
||||||
|
assert_eq!(json.get("memory.v1"), Some(&json!({"key": "value"})));
|
||||||
|
|
||||||
|
// Deserialize back
|
||||||
|
let deserialized: ExtensionData = serde_json::from_value(json).unwrap();
|
||||||
|
assert_eq!(
|
||||||
|
deserialized.get_extension_state("todo", "v0"),
|
||||||
|
Some(&json!("TODO content"))
|
||||||
|
);
|
||||||
|
assert_eq!(
|
||||||
|
deserialized.get_extension_state("memory", "v1"),
|
||||||
|
Some(&json!({"key": "value"}))
|
||||||
|
);
|
||||||
|
}
|
||||||
|
}
|
||||||
@@ -1,3 +1,4 @@
|
|||||||
|
pub mod extension_data;
|
||||||
pub mod info;
|
pub mod info;
|
||||||
pub mod storage;
|
pub mod storage;
|
||||||
|
|
||||||
@@ -9,4 +10,5 @@ pub use storage::{
|
|||||||
SessionMetadata,
|
SessionMetadata,
|
||||||
};
|
};
|
||||||
|
|
||||||
|
pub use extension_data::{ExtensionData, ExtensionState, TodoState};
|
||||||
pub use info::{get_valid_sorted_sessions, SessionInfo};
|
pub use info::{get_valid_sorted_sessions, SessionInfo};
|
||||||
|
|||||||
@@ -8,6 +8,7 @@
|
|||||||
use crate::conversation::message::Message;
|
use crate::conversation::message::Message;
|
||||||
use crate::conversation::Conversation;
|
use crate::conversation::Conversation;
|
||||||
use crate::providers::base::Provider;
|
use crate::providers::base::Provider;
|
||||||
|
use crate::session::extension_data::ExtensionData;
|
||||||
use crate::utils::safe_truncate;
|
use crate::utils::safe_truncate;
|
||||||
use anyhow::Result;
|
use anyhow::Result;
|
||||||
use chrono::Local;
|
use chrono::Local;
|
||||||
@@ -64,11 +65,13 @@ pub struct SessionMetadata {
|
|||||||
pub accumulated_input_tokens: Option<i32>,
|
pub accumulated_input_tokens: Option<i32>,
|
||||||
/// The number of output tokens used in the session. Accumulated across all messages.
|
/// The number of output tokens used in the session. Accumulated across all messages.
|
||||||
pub accumulated_output_tokens: Option<i32>,
|
pub accumulated_output_tokens: Option<i32>,
|
||||||
/// Session-scoped TODO list content
|
|
||||||
pub todo_content: Option<String>,
|
/// Extension data containing extension states
|
||||||
|
#[serde(default)]
|
||||||
|
pub extension_data: ExtensionData,
|
||||||
}
|
}
|
||||||
|
|
||||||
// Custom deserializer to handle old sessions without working_dir and todo_content
|
// Custom deserializer to handle old sessions without working_dir
|
||||||
impl<'de> Deserialize<'de> for SessionMetadata {
|
impl<'de> Deserialize<'de> for SessionMetadata {
|
||||||
fn deserialize<D>(deserializer: D) -> Result<Self, D::Error>
|
fn deserialize<D>(deserializer: D) -> Result<Self, D::Error>
|
||||||
where
|
where
|
||||||
@@ -78,7 +81,7 @@ impl<'de> Deserialize<'de> for SessionMetadata {
|
|||||||
struct Helper {
|
struct Helper {
|
||||||
description: String,
|
description: String,
|
||||||
message_count: usize,
|
message_count: usize,
|
||||||
schedule_id: Option<String>, // For backward compatibility
|
schedule_id: Option<String>,
|
||||||
total_tokens: Option<i32>,
|
total_tokens: Option<i32>,
|
||||||
input_tokens: Option<i32>,
|
input_tokens: Option<i32>,
|
||||||
output_tokens: Option<i32>,
|
output_tokens: Option<i32>,
|
||||||
@@ -86,7 +89,8 @@ impl<'de> Deserialize<'de> for SessionMetadata {
|
|||||||
accumulated_input_tokens: Option<i32>,
|
accumulated_input_tokens: Option<i32>,
|
||||||
accumulated_output_tokens: Option<i32>,
|
accumulated_output_tokens: Option<i32>,
|
||||||
working_dir: Option<PathBuf>,
|
working_dir: Option<PathBuf>,
|
||||||
todo_content: Option<String>, // For backward compatibility
|
#[serde(default)]
|
||||||
|
extension_data: ExtensionData,
|
||||||
}
|
}
|
||||||
|
|
||||||
let helper = Helper::deserialize(deserializer)?;
|
let helper = Helper::deserialize(deserializer)?;
|
||||||
@@ -108,7 +112,7 @@ impl<'de> Deserialize<'de> for SessionMetadata {
|
|||||||
accumulated_input_tokens: helper.accumulated_input_tokens,
|
accumulated_input_tokens: helper.accumulated_input_tokens,
|
||||||
accumulated_output_tokens: helper.accumulated_output_tokens,
|
accumulated_output_tokens: helper.accumulated_output_tokens,
|
||||||
working_dir,
|
working_dir,
|
||||||
todo_content: helper.todo_content,
|
extension_data: helper.extension_data,
|
||||||
})
|
})
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
@@ -133,7 +137,7 @@ impl SessionMetadata {
|
|||||||
accumulated_total_tokens: None,
|
accumulated_total_tokens: None,
|
||||||
accumulated_input_tokens: None,
|
accumulated_input_tokens: None,
|
||||||
accumulated_output_tokens: None,
|
accumulated_output_tokens: None,
|
||||||
todo_content: None,
|
extension_data: ExtensionData::new(),
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|||||||
@@ -411,6 +411,6 @@ pub fn create_test_session_metadata(message_count: usize, working_dir: &str) ->
|
|||||||
accumulated_total_tokens: Some(100),
|
accumulated_total_tokens: Some(100),
|
||||||
accumulated_input_tokens: Some(50),
|
accumulated_input_tokens: Some(50),
|
||||||
accumulated_output_tokens: Some(50),
|
accumulated_output_tokens: Some(50),
|
||||||
todo_content: None,
|
extension_data: Default::default(),
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|||||||
@@ -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
|
// Since we're using a mock provider, we can't test the actual TODO content
|
||||||
// but we can verify the metadata structure is correct
|
// 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]
|
#[tokio::test]
|
||||||
@@ -172,7 +175,11 @@ async fn test_todo_list_reads_from_session() {
|
|||||||
// Pre-populate session with TODO content
|
// Pre-populate session with TODO content
|
||||||
let session_path = goose::session::storage::get_path(session_id.clone()).unwrap();
|
let session_path = goose::session::storage::get_path(session_id.clone()).unwrap();
|
||||||
let mut metadata = SessionMetadata::default();
|
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)
|
goose::session::storage::update_metadata(&session_path, &metadata)
|
||||||
.await
|
.await
|
||||||
.unwrap();
|
.unwrap();
|
||||||
@@ -206,21 +213,27 @@ async fn test_todo_list_reads_from_session() {
|
|||||||
|
|
||||||
// Verify the TODO content is still in session
|
// Verify the TODO content is still in session
|
||||||
let metadata_after = goose::session::storage::read_metadata(&session_path).unwrap();
|
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!(
|
assert_eq!(
|
||||||
metadata_after.todo_content,
|
todo_state_after.unwrap().content,
|
||||||
Some("- Task 1\n- Task 2\n- Task 3".to_string())
|
"- Task 1\n- Task 2\n- Task 3".to_string()
|
||||||
);
|
);
|
||||||
}
|
}
|
||||||
|
|
||||||
#[tokio::test]
|
#[tokio::test]
|
||||||
async fn test_todo_isolation_between_sessions() {
|
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 session1_id = session::Identifier::Name(format!("test_session_{}", Uuid::new_v4()));
|
||||||
let session2_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
|
// Add TODO to session1
|
||||||
let session1_path = goose::session::storage::get_path(session1_id.clone()).unwrap();
|
let session1_path = goose::session::storage::get_path(session1_id.clone()).unwrap();
|
||||||
let mut metadata1 = SessionMetadata::default();
|
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)
|
goose::session::storage::update_metadata(&session1_path, &metadata1)
|
||||||
.await
|
.await
|
||||||
.unwrap();
|
.unwrap();
|
||||||
@@ -228,7 +241,10 @@ async fn test_todo_isolation_between_sessions() {
|
|||||||
// Add different TODO to session2
|
// Add different TODO to session2
|
||||||
let session2_path = goose::session::storage::get_path(session2_id.clone()).unwrap();
|
let session2_path = goose::session::storage::get_path(session2_id.clone()).unwrap();
|
||||||
let mut metadata2 = SessionMetadata::default();
|
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)
|
goose::session::storage::update_metadata(&session2_path, &metadata2)
|
||||||
.await
|
.await
|
||||||
.unwrap();
|
.unwrap();
|
||||||
@@ -237,12 +253,16 @@ async fn test_todo_isolation_between_sessions() {
|
|||||||
let metadata1_read = goose::session::storage::read_metadata(&session1_path).unwrap();
|
let metadata1_read = goose::session::storage::read_metadata(&session1_path).unwrap();
|
||||||
let metadata2_read = goose::session::storage::read_metadata(&session2_path).unwrap();
|
let metadata2_read = goose::session::storage::read_metadata(&session2_path).unwrap();
|
||||||
|
|
||||||
assert_eq!(metadata1_read.todo_content.unwrap(), "Session 1 tasks");
|
let todo1 = TodoState::from_extension_data(&metadata1_read.extension_data).unwrap();
|
||||||
assert_eq!(metadata2_read.todo_content.unwrap(), "Session 2 tasks");
|
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]
|
#[tokio::test]
|
||||||
async fn test_todo_clear_removes_from_session() {
|
async fn test_todo_clear_removes_from_session() {
|
||||||
|
use goose::session::extension_data::{ExtensionState, TodoState};
|
||||||
let temp_dir = create_test_session_dir().await;
|
let temp_dir = create_test_session_dir().await;
|
||||||
let session_id = session::Identifier::Name(format!("test_session_{}", Uuid::new_v4()));
|
let session_id = session::Identifier::Name(format!("test_session_{}", Uuid::new_v4()));
|
||||||
let agent = create_test_agent_with_mock_provider().await;
|
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
|
// Pre-populate session with TODO content
|
||||||
let session_path = goose::session::storage::get_path(session_id.clone()).unwrap();
|
let session_path = goose::session::storage::get_path(session_id.clone()).unwrap();
|
||||||
let mut metadata = SessionMetadata::default();
|
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)
|
goose::session::storage::update_metadata(&session_path, &metadata)
|
||||||
.await
|
.await
|
||||||
.unwrap();
|
.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
|
// With mock provider, the TODO won't actually be cleared via tool calls
|
||||||
// but we can verify the structure is correct
|
// but we can verify the structure is correct
|
||||||
let metadata_after = goose::session::storage::read_metadata(&session_path).unwrap();
|
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]
|
#[tokio::test]
|
||||||
async fn test_todo_persistence_across_agent_instances() {
|
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()));
|
let session_id = session::Identifier::Name(format!("test_session_{}", Uuid::new_v4()));
|
||||||
|
|
||||||
// First agent instance adds TODO
|
// First agent instance adds TODO
|
||||||
{
|
{
|
||||||
let session_path = goose::session::storage::get_path(session_id.clone()).unwrap();
|
let session_path = goose::session::storage::get_path(session_id.clone()).unwrap();
|
||||||
let mut metadata = SessionMetadata::default();
|
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)
|
goose::session::storage::update_metadata(&session_path, &metadata)
|
||||||
.await
|
.await
|
||||||
.unwrap();
|
.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 session_path = goose::session::storage::get_path(session_id.clone()).unwrap();
|
||||||
let metadata = goose::session::storage::read_metadata(&session_path).unwrap();
|
let metadata = goose::session::storage::read_metadata(&session_path).unwrap();
|
||||||
|
let todo_state = TodoState::from_extension_data(&metadata.extension_data).unwrap();
|
||||||
assert_eq!(metadata.todo_content.unwrap(), "Persistent task");
|
assert_eq!(todo_state.content, "Persistent task");
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
#[tokio::test]
|
#[tokio::test]
|
||||||
async fn test_todo_max_chars_limit() {
|
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()));
|
let session_id = session::Identifier::Name(format!("test_session_{}", Uuid::new_v4()));
|
||||||
|
|
||||||
// Set a small limit for testing
|
// 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
|
// Try to set content that exceeds the limit
|
||||||
let long_content = "x".repeat(100);
|
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)
|
// This should succeed at the storage level (storage doesn't enforce limits)
|
||||||
goose::session::storage::update_metadata(&session_path, &metadata)
|
goose::session::storage::update_metadata(&session_path, &metadata)
|
||||||
@@ -334,6 +366,7 @@ async fn test_todo_max_chars_limit() {
|
|||||||
|
|
||||||
#[tokio::test]
|
#[tokio::test]
|
||||||
async fn test_todo_with_special_characters() {
|
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_id = session::Identifier::Name(format!("test_session_{}", Uuid::new_v4()));
|
||||||
|
|
||||||
let session_path = goose::session::storage::get_path(session_id.clone()).unwrap();
|
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
|
- 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)
|
goose::session::storage::update_metadata(&session_path, &metadata)
|
||||||
.await
|
.await
|
||||||
.unwrap();
|
.unwrap();
|
||||||
|
|
||||||
// Read back and verify
|
// Read back and verify
|
||||||
let metadata_read = goose::session::storage::read_metadata(&session_path).unwrap();
|
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]
|
#[tokio::test]
|
||||||
async fn test_todo_concurrent_access() {
|
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()));
|
let session_id = session::Identifier::Name(format!("test_session_{}", Uuid::new_v4()));
|
||||||
|
|
||||||
// Spawn multiple concurrent TODO operations
|
// 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)
|
let mut metadata = goose::session::storage::read_metadata(&session_path)
|
||||||
.unwrap_or_else(|_| SessionMetadata::default());
|
.unwrap_or_else(|_| SessionMetadata::default());
|
||||||
|
|
||||||
let current_content = metadata.todo_content.unwrap_or_default();
|
let current_content = TodoState::from_extension_data(&metadata.extension_data)
|
||||||
metadata.todo_content = Some(format!("{}\n- Task {}", current_content, i));
|
.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
|
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
|
// Verify final state contains at least one task
|
||||||
let session_path = goose::session::storage::get_path(session_id).unwrap();
|
let session_path = goose::session::storage::get_path(session_id).unwrap();
|
||||||
let metadata = goose::session::storage::read_metadata(&session_path).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)
|
// Should contain at least one task (concurrent writes may overwrite)
|
||||||
assert!(todo_content.contains("Task"));
|
assert!(todo_state.content.contains("Task"));
|
||||||
}
|
}
|
||||||
|
|
||||||
#[tokio::test]
|
#[tokio::test]
|
||||||
async fn test_todo_empty_session_returns_empty() {
|
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_id = session::Identifier::Name(format!("test_session_{}", Uuid::new_v4()));
|
||||||
|
|
||||||
let session_path = goose::session::storage::get_path(session_id.clone()).unwrap();
|
let session_path = goose::session::storage::get_path(session_id.clone()).unwrap();
|
||||||
let metadata = goose::session::storage::read_metadata(&session_path)
|
let metadata = goose::session::storage::read_metadata(&session_path)
|
||||||
.unwrap_or_else(|_| SessionMetadata::default());
|
.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]
|
#[tokio::test]
|
||||||
async fn test_todo_update_preserves_other_metadata() {
|
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_id = session::Identifier::Name(format!("test_session_{}", Uuid::new_v4()));
|
||||||
|
|
||||||
let session_path = goose::session::storage::get_path(session_id.clone()).unwrap();
|
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.message_count = 5;
|
||||||
metadata.description = "Test session".to_string();
|
metadata.description = "Test session".to_string();
|
||||||
metadata.total_tokens = Some(1000);
|
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)
|
goose::session::storage::update_metadata(&session_path, &metadata)
|
||||||
.await
|
.await
|
||||||
.unwrap();
|
.unwrap();
|
||||||
|
|
||||||
// Update only TODO content
|
// 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)
|
goose::session::storage::update_metadata(&session_path, &metadata)
|
||||||
.await
|
.await
|
||||||
.unwrap();
|
.unwrap();
|
||||||
@@ -437,5 +489,6 @@ async fn test_todo_update_preserves_other_metadata() {
|
|||||||
assert_eq!(metadata_read.message_count, 5);
|
assert_eq!(metadata_read.message_count, 5);
|
||||||
assert_eq!(metadata_read.description, "Test session");
|
assert_eq!(metadata_read.description, "Test session");
|
||||||
assert_eq!(metadata_read.total_tokens, Some(1000));
|
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");
|
||||||
}
|
}
|
||||||
|
|||||||
@@ -2123,6 +2123,11 @@
|
|||||||
"propertyName": "type"
|
"propertyName": "type"
|
||||||
}
|
}
|
||||||
},
|
},
|
||||||
|
"ExtensionData": {
|
||||||
|
"type": "object",
|
||||||
|
"description": "Extension data containing all extension states\nKeys are in format \"extension_name.version\" (e.g., \"todo.v0\")",
|
||||||
|
"additionalProperties": {}
|
||||||
|
},
|
||||||
"ExtensionEntry": {
|
"ExtensionEntry": {
|
||||||
"allOf": [
|
"allOf": [
|
||||||
{
|
{
|
||||||
@@ -3207,6 +3212,9 @@
|
|||||||
"type": "string",
|
"type": "string",
|
||||||
"description": "A short description of the session, typically 3 words or less"
|
"description": "A short description of the session, typically 3 words or less"
|
||||||
},
|
},
|
||||||
|
"extension_data": {
|
||||||
|
"$ref": "#/components/schemas/ExtensionData"
|
||||||
|
},
|
||||||
"input_tokens": {
|
"input_tokens": {
|
||||||
"type": "integer",
|
"type": "integer",
|
||||||
"format": "int32",
|
"format": "int32",
|
||||||
@@ -3229,11 +3237,6 @@
|
|||||||
"description": "ID of the schedule that triggered this session, if any",
|
"description": "ID of the schedule that triggered this session, if any",
|
||||||
"nullable": true
|
"nullable": true
|
||||||
},
|
},
|
||||||
"todo_content": {
|
|
||||||
"type": "string",
|
|
||||||
"description": "Session-scoped TODO list content",
|
|
||||||
"nullable": true
|
|
||||||
},
|
|
||||||
"total_tokens": {
|
"total_tokens": {
|
||||||
"type": "integer",
|
"type": "integer",
|
||||||
"format": "int32",
|
"format": "int32",
|
||||||
|
|||||||
@@ -279,6 +279,14 @@ export type ExtensionConfig = {
|
|||||||
type: 'inline_python';
|
type: 'inline_python';
|
||||||
};
|
};
|
||||||
|
|
||||||
|
/**
|
||||||
|
* Extension data containing all extension states
|
||||||
|
* Keys are in format "extension_name.version" (e.g., "todo.v0")
|
||||||
|
*/
|
||||||
|
export type ExtensionData = {
|
||||||
|
[key: string]: unknown;
|
||||||
|
};
|
||||||
|
|
||||||
export type ExtensionEntry = ExtensionConfig & {
|
export type ExtensionEntry = ExtensionConfig & {
|
||||||
type?: 'ExtensionEntry';
|
type?: 'ExtensionEntry';
|
||||||
} & {
|
} & {
|
||||||
@@ -682,6 +690,7 @@ export type SessionMetadata = {
|
|||||||
* A short description of the session, typically 3 words or less
|
* A short description of the session, typically 3 words or less
|
||||||
*/
|
*/
|
||||||
description: string;
|
description: string;
|
||||||
|
extension_data?: ExtensionData;
|
||||||
/**
|
/**
|
||||||
* The number of input tokens used in the session. Retrieved from the provider's last usage.
|
* The number of input tokens used in the session. Retrieved from the provider's last usage.
|
||||||
*/
|
*/
|
||||||
@@ -698,10 +707,6 @@ export type SessionMetadata = {
|
|||||||
* ID of the schedule that triggered this session, if any
|
* ID of the schedule that triggered this session, if any
|
||||||
*/
|
*/
|
||||||
schedule_id?: string | null;
|
schedule_id?: string | null;
|
||||||
/**
|
|
||||||
* Session-scoped TODO list content
|
|
||||||
*/
|
|
||||||
todo_content?: string | null;
|
|
||||||
/**
|
/**
|
||||||
* The total number of tokens used in the session. Retrieved from the provider's last usage.
|
* The total number of tokens used in the session. Retrieved from the provider's last usage.
|
||||||
*/
|
*/
|
||||||
|
|||||||
Reference in New Issue
Block a user