diff --git a/src/gateway/mod.rs b/src/gateway/mod.rs index 69cdf0e..6f942e2 100644 --- a/src/gateway/mod.rs +++ b/src/gateway/mod.rs @@ -43,6 +43,7 @@ impl GatewayState { session_ttl_hours, agent_prompt_reinject_every, show_tool_results, + config.time.timezone.clone(), provider_config, provider_configs, skills, diff --git a/src/gateway/session.rs b/src/gateway/session.rs index 66cba6c..0d2b0cf 100644 --- a/src/gateway/session.rs +++ b/src/gateway/session.rs @@ -23,7 +23,7 @@ use crate::storage::{SessionRecord, SessionStore, persistent_session_id}; use crate::tools::{ BashTool, CalculatorTool, FileEditTool, FileReadTool, FileWriteTool, HttpRequestTool, MemoryManageTool, MemorySearchTool, SchedulerManageTool, SkillListTool, SkillManageTool, ToolContext, ToolRegistry, - WebFetchTool, + TimeTool, WebFetchTool, }; const DEFAULT_AGENT_PROMPT: &str = include_str!("default_agent_prompt.md"); @@ -822,9 +822,15 @@ struct SessionManagerInner { session_ttl: Duration, } -fn default_tools(skills: Arc, store: Arc, known_agents: HashSet) -> ToolRegistry { +fn default_tools( + skills: Arc, + store: Arc, + known_agents: HashSet, + default_timezone: String, +) -> ToolRegistry { let mut registry = ToolRegistry::new(); registry.register(CalculatorTool::new()); + registry.register(TimeTool::new(default_timezone)); registry.register(FileReadTool::new()); registry.register(FileWriteTool::new()); registry.register(FileEditTool::new()); @@ -963,6 +969,7 @@ impl SessionManager { session_ttl_hours: u64, agent_prompt_reinject_every: u64, show_tool_results: bool, + default_timezone: String, provider_config: LLMProviderConfig, provider_configs: HashMap, skills: Arc, @@ -985,7 +992,12 @@ impl SessionManager { })), provider_config, provider_configs: Arc::new(provider_configs), - tools: Arc::new(default_tools(skills.clone(), store.clone(), known_agents)), + tools: Arc::new(default_tools( + skills.clone(), + store.clone(), + known_agents, + default_timezone, + )), skills, store, agent_prompt_reinject_every, @@ -1690,7 +1702,12 @@ mod tests { let store = Arc::new(SessionStore::in_memory().unwrap()); let (user_tx, _user_rx) = mpsc::channel(4); let skills = Arc::new(SkillRuntime::default()); - let tools = Arc::new(default_tools(skills.clone(), store.clone(), HashSet::new())); + let tools = Arc::new(default_tools( + skills.clone(), + store.clone(), + HashSet::new(), + "Asia/Shanghai".to_string(), + )); let mut session = Session::new( "feishu".to_string(), test_provider_config(), @@ -1863,6 +1880,7 @@ mod tests { 4, 100, false, + "Asia/Shanghai".to_string(), provider_config.clone(), HashMap::from([("default".to_string(), provider_config)]), Arc::new(SkillRuntime::default()), @@ -1906,6 +1924,7 @@ mod tests { 4, 100, false, + "Asia/Shanghai".to_string(), default_provider.clone(), HashMap::from([ ("default".to_string(), default_provider), @@ -1971,6 +1990,7 @@ mod tests { 4, 100, false, + "Asia/Shanghai".to_string(), provider_config.clone(), HashMap::from([("default".to_string(), provider_config)]), Arc::new(SkillRuntime::default()), @@ -2045,6 +2065,7 @@ mod tests { 4, 100, false, + "Asia/Shanghai".to_string(), provider_config.clone(), HashMap::from([("default".to_string(), provider_config)]), Arc::new(SkillRuntime::default()), @@ -2134,6 +2155,7 @@ mod tests { 4, 100, false, + "Asia/Shanghai".to_string(), provider_config.clone(), HashMap::from([("default".to_string(), provider_config)]), Arc::new(SkillRuntime::default()), @@ -2194,6 +2216,7 @@ mod tests { 4, 100, false, + "Asia/Shanghai".to_string(), provider_config.clone(), HashMap::from([("default".to_string(), provider_config)]), Arc::new(SkillRuntime::default()), @@ -2364,7 +2387,12 @@ mod tests { let store = Arc::new(SessionStore::in_memory().unwrap()); let (user_tx, _user_rx) = mpsc::channel(4); let skills = Arc::new(SkillRuntime::default()); - let tools = Arc::new(default_tools(skills.clone(), store.clone(), HashSet::new())); + let tools = Arc::new(default_tools( + skills.clone(), + store.clone(), + HashSet::new(), + "Asia/Shanghai".to_string(), + )); let mut session = Session::new( "feishu".to_string(), test_provider_config(), @@ -2412,7 +2440,12 @@ mod tests { let store = Arc::new(SessionStore::in_memory().unwrap()); let (user_tx, _user_rx) = mpsc::channel(4); let skills = Arc::new(SkillRuntime::default()); - let tools = Arc::new(default_tools(skills.clone(), store.clone(), HashSet::new())); + let tools = Arc::new(default_tools( + skills.clone(), + store.clone(), + HashSet::new(), + "Asia/Shanghai".to_string(), + )); let mut session = Session::new( "feishu".to_string(), test_provider_config(), @@ -2439,7 +2472,12 @@ mod tests { let store = Arc::new(SessionStore::in_memory().unwrap()); let (user_tx, _user_rx) = mpsc::channel(4); let skills = Arc::new(SkillRuntime::default()); - let tools = Arc::new(default_tools(skills.clone(), store.clone(), HashSet::new())); + let tools = Arc::new(default_tools( + skills.clone(), + store.clone(), + HashSet::new(), + "Asia/Shanghai".to_string(), + )); let mut session = Session::new( "feishu".to_string(), test_provider_config(), @@ -2484,7 +2522,12 @@ mod tests { let store = Arc::new(SessionStore::in_memory().unwrap()); let (user_tx, _user_rx) = mpsc::channel(4); let skills = Arc::new(SkillRuntime::default()); - let tools = Arc::new(default_tools(skills.clone(), store.clone(), HashSet::new())); + let tools = Arc::new(default_tools( + skills.clone(), + store.clone(), + HashSet::new(), + "Asia/Shanghai".to_string(), + )); let mut session = Session::new( "feishu".to_string(), test_provider_config(), @@ -2518,7 +2561,26 @@ mod tests { let store = Arc::new(SessionStore::in_memory().unwrap()); let (user_tx, _user_rx) = mpsc::channel(4); let skills = Arc::new(SkillRuntime::default()); - let tools = Arc::new(default_tools(skills.clone(), store.clone(), HashSet::new())); + let tools = Arc::new(default_tools( + skills.clone(), + store.clone(), + HashSet::new(), + "Asia/Shanghai".to_string(), + )); + + #[test] + fn test_default_tools_registers_get_time() { + let skills = Arc::new(SkillRuntime::default()); + let store = Arc::new(SessionStore::in_memory().unwrap()); + let tools = default_tools( + skills, + store, + HashSet::new(), + "Asia/Shanghai".to_string(), + ); + + assert!(tools.tool_names().iter().any(|name| name == "get_time")); + } let mut session = Session::new( "feishu".to_string(), test_provider_config(), diff --git a/src/scheduler/mod.rs b/src/scheduler/mod.rs index 6b634ef..e6627bf 100644 --- a/src/scheduler/mod.rs +++ b/src/scheduler/mod.rs @@ -856,6 +856,7 @@ mod tests { 4, 100, false, + "Asia/Shanghai".to_string(), provider_config.clone(), HashMap::from([("default".to_string(), provider_config)]), Arc::new(SkillRuntime::default()), @@ -906,6 +907,7 @@ mod tests { 4, 100, false, + "Asia/Shanghai".to_string(), provider_config.clone(), HashMap::from([("default".to_string(), provider_config)]), Arc::new(SkillRuntime::default()), diff --git a/src/tools/mod.rs b/src/tools/mod.rs index 53b29ef..3678fc7 100644 --- a/src/tools/mod.rs +++ b/src/tools/mod.rs @@ -10,6 +10,7 @@ pub mod registry; pub mod scheduler_manage; pub mod schema; pub mod skill_manage; +pub mod time; pub mod traits; pub mod web_fetch; @@ -25,5 +26,6 @@ pub use registry::ToolRegistry; pub use scheduler_manage::SchedulerManageTool; pub use schema::{CleaningStrategy, SchemaCleanr}; pub use skill_manage::{SkillListTool, SkillManageTool}; +pub use time::TimeTool; pub use traits::{Tool, ToolContext, ToolResult}; pub use web_fetch::WebFetchTool; diff --git a/src/tools/time.rs b/src/tools/time.rs new file mode 100644 index 0000000..a24f9ae --- /dev/null +++ b/src/tools/time.rs @@ -0,0 +1,447 @@ +use async_trait::async_trait; +use chrono::{DateTime, Days, Duration, Months, Utc}; +use serde_json::{Value, json}; + +use super::traits::{Tool, ToolResult}; + +pub struct TimeTool { + default_timezone: String, +} + +impl TimeTool { + pub fn new(default_timezone: impl Into) -> Self { + Self { + default_timezone: default_timezone.into(), + } + } + + fn execute_at(&self, now_utc: DateTime, args: Value) -> ToolResult { + match execute_time_request(now_utc, &self.default_timezone, args) { + Ok(output) => ToolResult { + success: true, + output, + error: None, + }, + Err(error) => ToolResult { + success: false, + output: String::new(), + error: Some(error), + }, + } + } +} + +impl Default for TimeTool { + fn default() -> Self { + Self::new("Asia/Shanghai") + } +} + +#[async_trait] +impl Tool for TimeTool { + fn name(&self) -> &str { + "get_time" + } + + fn description(&self) -> &str { + "Get the current time or calculate a time before/after now in minutes, hours, days, months, or years. Use this for current time, relative time offsets, or when exact timestamps matter. If amount/unit/direction are omitted, it returns the current time in the default timezone. Examples: current time, 30 minutes later, 2 hours ago, 7 days later, 1 month ago, 1 year later." + } + + fn read_only(&self) -> bool { + true + } + + fn parameters_schema(&self) -> serde_json::Value { + json!({ + "type": "object", + "properties": { + "direction": { + "type": "string", + "description": "Relative direction from the current time. Optional when only the current time is needed.", + "enum": ["past", "future", "before", "after", "ago", "later"] + }, + "amount": { + "type": "integer", + "description": "How many units to move relative to now. Required when direction and unit are used.", + "minimum": 0 + }, + "unit": { + "type": "string", + "description": "Time unit for the relative offset. Required when direction and amount are used.", + "enum": [ + "minute", "minutes", + "hour", "hours", + "day", "days", + "month", "months", + "year", "years" + ] + }, + "timezone": { + "type": "string", + "description": "Optional IANA timezone override, for example Asia/Shanghai or Europe/Berlin. Defaults to the PicoBot configured timezone." + } + } + }) + } + + async fn execute(&self, args: serde_json::Value) -> anyhow::Result { + Ok(self.execute_at(Utc::now(), args)) + } +} + +fn execute_time_request( + now_utc: DateTime, + default_timezone: &str, + args: Value, +) -> Result { + let timezone_name = args + .get("timezone") + .and_then(Value::as_str) + .unwrap_or(default_timezone); + let timezone = timezone_name.parse::().map_err(|_| { + format!( + "Invalid timezone: {timezone_name}. Expected an IANA timezone like Asia/Shanghai" + ) + })?; + + let now_local = now_utc.with_timezone(&timezone); + let offset = parse_offset_request(&args)?; + + let (kind, result_local, description, offset_json) = match offset { + None => ( + "current", + now_local, + format!( + "Current time in {timezone_name} is {}", + format_human_time(&now_local) + ), + Value::Null, + ), + Some(offset) => { + let result_local = apply_offset(now_local, &offset)?; + let direction_text = match offset.direction { + OffsetDirection::Past => "before", + OffsetDirection::Future => "after", + }; + let unit_label = pluralized_unit(offset.unit, offset.amount); + let description = format!( + "{} {} {} current time in {timezone_name} is {}", + offset.amount, + unit_label, + direction_text, + format_human_time(&result_local) + ); + ( + "offset", + result_local, + description, + json!({ + "direction": offset.direction.as_str(), + "amount": offset.amount, + "unit": offset.unit.as_str(), + }), + ) + } + }; + + serde_json::to_string_pretty(&json!({ + "kind": kind, + "timezone": timezone_name, + "current_time": now_local.to_rfc3339(), + "result_time": result_local.to_rfc3339(), + "unix_timestamp": result_local.timestamp(), + "offset": offset_json, + "description": description, + })) + .map_err(|error| format!("Failed to serialize time tool output: {error}")) +} + +fn parse_offset_request(args: &Value) -> Result, String> { + let direction = args.get("direction").and_then(Value::as_str); + let amount = args.get("amount"); + let unit = args.get("unit").and_then(Value::as_str); + + if direction.is_none() && amount.is_none() && unit.is_none() { + return Ok(None); + } + + let direction = direction.ok_or_else(|| { + "Missing required parameter: direction when requesting a relative time".to_string() + })?; + let amount = amount + .and_then(Value::as_u64) + .ok_or_else(|| "Missing required parameter: amount when requesting a relative time".to_string())?; + let amount = u32::try_from(amount) + .map_err(|_| "amount is too large; expected a 32-bit unsigned integer".to_string())?; + let unit = unit + .ok_or_else(|| "Missing required parameter: unit when requesting a relative time".to_string())?; + + Ok(Some(OffsetRequest { + direction: OffsetDirection::parse(direction)?, + amount, + unit: TimeUnit::parse(unit)?, + })) +} + +fn apply_offset( + now_local: DateTime, + offset: &OffsetRequest, +) -> Result, String> { + match (offset.direction, offset.unit) { + (OffsetDirection::Future, TimeUnit::Minute) => Ok(now_local + Duration::minutes(i64::from(offset.amount))), + (OffsetDirection::Past, TimeUnit::Minute) => Ok(now_local - Duration::minutes(i64::from(offset.amount))), + (OffsetDirection::Future, TimeUnit::Hour) => Ok(now_local + Duration::hours(i64::from(offset.amount))), + (OffsetDirection::Past, TimeUnit::Hour) => Ok(now_local - Duration::hours(i64::from(offset.amount))), + (OffsetDirection::Future, TimeUnit::Day) => now_local + .checked_add_days(Days::new(u64::from(offset.amount))) + .ok_or_else(|| "Failed to add days to the current time".to_string()), + (OffsetDirection::Past, TimeUnit::Day) => now_local + .checked_sub_days(Days::new(u64::from(offset.amount))) + .ok_or_else(|| "Failed to subtract days from the current time".to_string()), + (OffsetDirection::Future, TimeUnit::Month) => now_local + .checked_add_months(Months::new(offset.amount)) + .ok_or_else(|| "Failed to add months to the current time".to_string()), + (OffsetDirection::Past, TimeUnit::Month) => now_local + .checked_sub_months(Months::new(offset.amount)) + .ok_or_else(|| "Failed to subtract months from the current time".to_string()), + (OffsetDirection::Future, TimeUnit::Year) => now_local + .checked_add_months(Months::new(years_to_months(offset.amount)?)) + .ok_or_else(|| "Failed to add years to the current time".to_string()), + (OffsetDirection::Past, TimeUnit::Year) => now_local + .checked_sub_months(Months::new(years_to_months(offset.amount)?)) + .ok_or_else(|| "Failed to subtract years from the current time".to_string()), + } +} + +fn years_to_months(years: u32) -> Result { + years + .checked_mul(12) + .ok_or_else(|| "amount is too large to convert years into months".to_string()) +} + +fn format_human_time(value: &DateTime) -> String { + value.format("%Y-%m-%d %H:%M:%S %Z").to_string() +} + +fn pluralized_unit(unit: TimeUnit, amount: u32) -> &'static str { + match (unit, amount) { + (TimeUnit::Minute, 1) => "minute", + (TimeUnit::Minute, _) => "minutes", + (TimeUnit::Hour, 1) => "hour", + (TimeUnit::Hour, _) => "hours", + (TimeUnit::Day, 1) => "day", + (TimeUnit::Day, _) => "days", + (TimeUnit::Month, 1) => "month", + (TimeUnit::Month, _) => "months", + (TimeUnit::Year, 1) => "year", + (TimeUnit::Year, _) => "years", + } +} + +#[derive(Clone, Copy)] +struct OffsetRequest { + direction: OffsetDirection, + amount: u32, + unit: TimeUnit, +} + +#[derive(Clone, Copy)] +enum OffsetDirection { + Past, + Future, +} + +impl OffsetDirection { + fn parse(value: &str) -> Result { + match value.trim().to_ascii_lowercase().as_str() { + "past" | "before" | "ago" => Ok(Self::Past), + "future" | "after" | "later" => Ok(Self::Future), + other => Err(format!( + "Invalid direction: {other}. Expected one of: past, future, before, after, ago, later" + )), + } + } + + fn as_str(&self) -> &'static str { + match self { + Self::Past => "past", + Self::Future => "future", + } + } +} + +#[derive(Clone, Copy)] +enum TimeUnit { + Minute, + Hour, + Day, + Month, + Year, +} + +impl TimeUnit { + fn parse(value: &str) -> Result { + match value.trim().to_ascii_lowercase().as_str() { + "minute" | "minutes" => Ok(Self::Minute), + "hour" | "hours" => Ok(Self::Hour), + "day" | "days" => Ok(Self::Day), + "month" | "months" => Ok(Self::Month), + "year" | "years" => Ok(Self::Year), + other => Err(format!( + "Invalid unit: {other}. Expected one of: minute, hour, day, month, year" + )), + } + } + + fn as_str(&self) -> &'static str { + match self { + Self::Minute => "minute", + Self::Hour => "hour", + Self::Day => "day", + Self::Month => "month", + Self::Year => "year", + } + } +} + +#[cfg(test)] +mod tests { + use super::*; + use chrono::{TimeZone, Timelike}; + + fn fixed_now() -> DateTime { + Utc.with_ymd_and_hms(2026, 4, 27, 4, 30, 0) + .single() + .unwrap() + } + + #[test] + fn test_get_current_time_uses_default_timezone() { + let tool = TimeTool::new("Asia/Shanghai"); + let result = tool.execute_at(fixed_now(), json!({})); + + assert!(result.success); + let payload: Value = serde_json::from_str(&result.output).unwrap(); + assert_eq!(payload["kind"], "current"); + assert_eq!(payload["timezone"], "Asia/Shanghai"); + assert_eq!(payload["result_time"], "2026-04-27T12:30:00+08:00"); + } + + #[test] + fn test_get_future_hours() { + let tool = TimeTool::new("Asia/Shanghai"); + let result = tool.execute_at( + fixed_now(), + json!({"direction": "future", "amount": 2, "unit": "hours"}), + ); + + assert!(result.success); + let payload: Value = serde_json::from_str(&result.output).unwrap(); + assert_eq!(payload["result_time"], "2026-04-27T14:30:00+08:00"); + assert_eq!(payload["offset"]["direction"], "future"); + assert_eq!(payload["offset"]["unit"], "hour"); + } + + #[test] + fn test_get_past_minutes() { + let tool = TimeTool::new("Asia/Shanghai"); + let result = tool.execute_at( + fixed_now(), + json!({"direction": "ago", "amount": 15, "unit": "minute"}), + ); + + assert!(result.success); + let payload: Value = serde_json::from_str(&result.output).unwrap(); + assert_eq!(payload["result_time"], "2026-04-27T12:15:00+08:00"); + } + + #[test] + fn test_get_future_days_crosses_date_boundary() { + let tool = TimeTool::new("Asia/Shanghai"); + let result = tool.execute_at( + fixed_now(), + json!({"direction": "after", "amount": 3, "unit": "days"}), + ); + + assert!(result.success); + let payload: Value = serde_json::from_str(&result.output).unwrap(); + assert_eq!(payload["result_time"], "2026-04-30T12:30:00+08:00"); + } + + #[test] + fn test_get_past_month_rolls_back_to_last_valid_day() { + let tool = TimeTool::new("Asia/Shanghai"); + let now = Utc.with_ymd_and_hms(2026, 3, 31, 1, 0, 0).single().unwrap(); + let result = tool.execute_at( + now, + json!({"direction": "past", "amount": 1, "unit": "month"}), + ); + + assert!(result.success); + let payload: Value = serde_json::from_str(&result.output).unwrap(); + assert_eq!(payload["result_time"], "2026-02-28T09:00:00+08:00"); + } + + #[test] + fn test_get_future_year_handles_leap_day() { + let tool = TimeTool::new("Asia/Shanghai"); + let now = Utc.with_ymd_and_hms(2024, 2, 29, 4, 0, 0).single().unwrap(); + let result = tool.execute_at( + now, + json!({"direction": "future", "amount": 1, "unit": "year"}), + ); + + assert!(result.success); + let payload: Value = serde_json::from_str(&result.output).unwrap(); + assert_eq!(payload["result_time"], "2025-02-28T12:00:00+08:00"); + } + + #[test] + fn test_timezone_override_is_used() { + let tool = TimeTool::new("Asia/Shanghai"); + let result = tool.execute_at(fixed_now(), json!({"timezone": "Europe/Berlin"})); + + assert!(result.success); + let payload: Value = serde_json::from_str(&result.output).unwrap(); + assert_eq!(payload["timezone"], "Europe/Berlin"); + let result_time = payload["result_time"].as_str().unwrap(); + assert!(result_time.ends_with("+02:00")); + } + + #[test] + fn test_missing_relative_unit_returns_error() { + let tool = TimeTool::new("Asia/Shanghai"); + let result = tool.execute_at(fixed_now(), json!({"direction": "future", "amount": 1})); + + assert!(!result.success); + assert_eq!( + result.error.as_deref(), + Some("Missing required parameter: unit when requesting a relative time") + ); + } + + #[test] + fn test_invalid_timezone_returns_error() { + let tool = TimeTool::new("Asia/Shanghai"); + let result = tool.execute_at(fixed_now(), json!({"timezone": "Mars/Base"})); + + assert!(!result.success); + assert!(result.error.unwrap().contains("Invalid timezone")); + } + + #[test] + fn test_zero_offset_keeps_same_local_time() { + let tool = TimeTool::new("Asia/Shanghai"); + let result = tool.execute_at( + fixed_now(), + json!({"direction": "future", "amount": 0, "unit": "hours"}), + ); + + assert!(result.success); + let payload: Value = serde_json::from_str(&result.output).unwrap(); + let result_time = chrono::DateTime::parse_from_rfc3339(payload["result_time"].as_str().unwrap()) + .unwrap(); + assert_eq!(result_time.hour(), 12); + assert_eq!(result_time.minute(), 30); + } +} \ No newline at end of file