Merge pull request #785 from dcSpark/nico/update_app_id

continuous app_id
dcSpark · Jan 10, 2025 · f842a60 · f842a60
2 parents 8ae193f + 6ab607b
commit f842a60
Show file tree

Hide file tree

Showing 3 changed files with 37 additions and 22 deletions.
diff --git a/shinkai-bin/shinkai-node/src/cron_tasks/cron_manager.rs b/shinkai-bin/shinkai-node/src/cron_tasks/cron_manager.rs
@@ -15,7 +15,7 @@ use shinkai_message_primitives::{
         shinkai_name::ShinkaiName,
         ws_types::WSUpdateHandler,
     },
-    shinkai_message::shinkai_message_schemas::JobMessage,
+    shinkai_message::shinkai_message_schemas::{AssociatedUI, JobMessage},
     shinkai_utils::{
         shinkai_logging::{shinkai_log, ShinkaiLogLevel, ShinkaiLogOption},
         signatures::clone_signature_secret_key,
@@ -331,6 +331,8 @@ impl CronManager {
                 if job_creation_info_clone.is_hidden.is_none() {
                     job_creation_info_clone.is_hidden = Some(true);
                 }
+                // Set the associated UI to Cron with the task ID
+                job_creation_info_clone.associated_ui = Some(AssociatedUI::Cron(cron_job.task_id.to_string()));
 
                 let job_id = job_manager
                     .lock()

diff --git a/shinkai-bin/shinkai-node/src/managers/tool_router.rs b/shinkai-bin/shinkai-node/src/managers/tool_router.rs
@@ -24,7 +24,7 @@ use shinkai_message_primitives::schemas::shinkai_tool_offering::{
 use shinkai_message_primitives::schemas::shinkai_tools::CodeLanguage;
 use shinkai_message_primitives::schemas::wallet_mixed::{Asset, NetworkIdentifier};
 use shinkai_message_primitives::schemas::ws_types::{PaymentMetadata, WSMessageType, WidgetMetadata};
-use shinkai_message_primitives::shinkai_message::shinkai_message_schemas::WSTopic;
+use shinkai_message_primitives::shinkai_message::shinkai_message_schemas::{AssociatedUI, WSTopic};
 use shinkai_message_primitives::shinkai_utils::shinkai_logging::{shinkai_log, ShinkaiLogLevel, ShinkaiLogOption};
 use shinkai_sqlite::errors::SqliteManagerError;
 use shinkai_sqlite::files::prompts_data;
@@ -166,13 +166,7 @@ impl ToolRouter {
 
         let tool_urls = tools
             .iter()
-            .map(|tool| {
-                (
-                    tool["name"].as_str(),
-                    tool["file"].as_str(),
-                    tool["router_key"].as_str(),
-                )
-            })
+            .map(|tool| (tool["name"].as_str(), tool["file"].as_str(), tool["routerKey"].as_str()))
             .collect::<Vec<_>>()
             .into_iter()
             .filter(|(name, url, router_key)| url.is_some() && name.is_some() && router_key.is_some())
@@ -496,7 +490,13 @@ impl ToolRouter {
                     .node_storage_path
                     .clone()
                     .ok_or_else(|| ToolError::ExecutionError("Node storage path is not set".to_string()))?;
-                let app_id = context.full_job().job_id().to_string();
+
+                // Get app_id from Cron UI if present, otherwise use job_id
+                let app_id = match context.full_job().associated_ui().as_ref() {
+                    Some(AssociatedUI::Cron(cron_id)) => cron_id.clone(),
+                    _ => context.full_job().job_id().to_string(),
+                };
+
                 let tool_id = shinkai_tool.tool_router_key().to_string_without_version().clone();
                 let tools = python_tool.tools.clone();
                 let support_files =
@@ -507,10 +507,10 @@ impl ToolRouter {
                 let envs = generate_execution_environment(
                     context.db(),
                     context.agent().clone().get_id().to_string(),
-                    format!("jid-{}", tool_id),
-                    format!("jid-{}", app_id),
+                    tool_id.clone(),
+                    app_id.clone(),
                     shinkai_tool.tool_router_key().to_string_without_version().clone(),
-                    format!("jid-{}", app_id),
+                    app_id.clone(),
                     &python_tool.oauth,
                 )
                 .await
@@ -539,8 +539,8 @@ impl ToolRouter {
                         function_args,
                         function_config_vec,
                         node_storage_path,
-                        app_id,
-                        tool_id,
+                        app_id.clone(),
+                        tool_id.clone(),
                         node_name,
                         false,
                         None,
@@ -555,7 +555,12 @@ impl ToolRouter {
                 });
             }
             ShinkaiTool::Rust(_rust_tool, _is_enabled) => {
-                let app_id = context.full_job().job_id().to_string();
+                // Get app_id from Cron UI if present, otherwise use job_id
+                let app_id = match context.full_job().associated_ui().as_ref() {
+                    Some(AssociatedUI::Cron(cron_id)) => cron_id.clone(),
+                    _ => context.full_job().job_id().to_string(),
+                };
+
                 let tool_id = shinkai_tool.tool_router_key().to_string_without_version().clone();
                 let function_config = shinkai_tool.get_config_from_env();
                 let function_config_vec: Vec<ToolConfig> = function_config.into_iter().collect();
@@ -564,7 +569,6 @@ impl ToolRouter {
                 let llm_provider = context.agent().get_llm_provider_id().to_string();
                 let bearer = db.read_api_v2_key().unwrap_or_default().unwrap_or_default();
 
-
                 let job_callback_manager = context.job_callback_manager();
                 let mut job_manager: Option<Arc<Mutex<JobManager>>> = None;
                 if let Some(job_callback_manager) = &job_callback_manager {
@@ -573,7 +577,9 @@ impl ToolRouter {
                 }
 
                 if job_manager.is_none() {
-                    return Err(LLMProviderError::FunctionExecutionError("Job manager is not available".to_string()));
+                    return Err(LLMProviderError::FunctionExecutionError(
+                        "Job manager is not available".to_string(),
+                    ));
                 }
 
                 let result = execute_custom_tool(
@@ -610,7 +616,13 @@ impl ToolRouter {
                     .node_storage_path
                     .clone()
                     .ok_or_else(|| ToolError::ExecutionError("Node storage path is not set".to_string()))?;
-                let app_id = context.full_job().job_id().to_string();
+
+                // Get app_id from Cron UI if present, otherwise use job_id
+                let app_id = match context.full_job().associated_ui().as_ref() {
+                    Some(AssociatedUI::Cron(cron_id)) => cron_id.clone(),
+                    _ => context.full_job().job_id().to_string(),
+                };
+
                 let tool_id = shinkai_tool.tool_router_key().to_string_without_version().clone();
                 let tools = deno_tool.tools.clone();
                 let support_files =
@@ -621,10 +633,10 @@ impl ToolRouter {
                 let envs = generate_execution_environment(
                     context.db(),
                     context.agent().clone().get_id().to_string(),
-                    format!("jid-{}", app_id),
-                    format!("jid-{}", tool_id),
+                    app_id.clone(),
+                    tool_id.clone(),
                     shinkai_tool.tool_router_key().to_string_without_version().clone(),
-                    format!("jid-{}", app_id),
+                    app_id.clone(),
                     &deno_tool.oauth,
                 )
                 .await

diff --git a/shinkai-libs/shinkai-message-primitives/src/shinkai_message/shinkai_message_schemas.rs b/shinkai-libs/shinkai-message-primitives/src/shinkai_message/shinkai_message_schemas.rs
@@ -296,6 +296,7 @@ pub struct SymmetricKeyExchange {
 pub enum AssociatedUI {
     Sheet(String),
     Playground,
+    Cron(String),
     // Add more variants as needed
 }