- fix: several fixes in shinkai node manager (#376)

* - fix: scape control characters from string in logs * - fix: api port * - fix: ollama and shinkai-node health proof * - fix: isInUse state when auth change and node is not running
dcSpark · Jul 18, 2024 · 0bc1228 · 0bc1228
1 parent 167692e
commit 0bc1228
Show file tree

Hide file tree

Showing 8 changed files with 89 additions and 53 deletions.
diff --git a/apps/shinkai-desktop/src-tauri/Cargo.lock b/apps/shinkai-desktop/src-tauri/Cargo.lock
diff --git a/apps/shinkai-desktop/src-tauri/Cargo.toml b/apps/shinkai-desktop/src-tauri/Cargo.toml
@@ -23,7 +23,7 @@ hound = "3.4.0"
 # fix this dependency later on
 reqwest = { version = "0.11", features = ["json", "stream"] }
 lazy_static = "1.4.0"
-tokio = "1.36.0"
+tokio = { version = "1.36.0", features = ["macros"] }
 chrono = "0.4.38"
 futures-util = "0.3"
 regex = "1.10.4"

diff --git a/apps/shinkai-desktop/src-tauri/src/local_shinkai_node/process_handlers/logger.rs b/apps/shinkai-desktop/src-tauri/src/local_shinkai_node/process_handlers/logger.rs
@@ -27,7 +27,7 @@ impl Logger {
         LogEntry {
             timestamp: current_timestamp,
             process: self.process_name.clone(),
-            message,
+            message: message.chars().filter(|&c| !c.is_control()).collect(),
         }
     }
 

diff --git a/...inkai-desktop/src-tauri/src/local_shinkai_node/process_handlers/ollama_process_handler.rs b/...inkai-desktop/src-tauri/src/local_shinkai_node/process_handlers/ollama_process_handler.rs
@@ -1,3 +1,5 @@
+use std::time::Duration;
+
 use regex::Regex;
 use serde::Serialize;
 use tokio::sync::mpsc::Sender;
@@ -24,7 +26,7 @@ pub struct OllamaProcessHandler {
 }
 
 impl OllamaProcessHandler {
-    const HEALTH_TIMEOUT_MS: u64 = 500;
+    const HEALTH_TIMEOUT_MS: u64 = 5000;
     const PROCESS_NAME: &'static str = "ollama";
     const READY_MATCHER: &'static str = "Listening on ";
 
@@ -58,23 +60,25 @@ impl OllamaProcessHandler {
     }
 
     async fn wait_ollama_server(&self) -> Result<(), String> {
+        let timeout = Duration::from_millis(Self::HEALTH_TIMEOUT_MS);
         let start_time = std::time::Instant::now();
-        let mut success = false;
         let ollama_api = OllamaApiClient::new(self.get_ollama_api_base_url());
-        while std::time::Instant::now().duration_since(start_time)
-            < std::time::Duration::from_millis(Self::HEALTH_TIMEOUT_MS)
-        {
-            let status = ollama_api.health().await;
-            if status.is_ok() && status.unwrap() {
-                success = true;
-                break;
+        tokio::select! {
+            _ = tokio::time::sleep(timeout) => {
+                let elapsed = start_time.elapsed();
+                Err(format!("wait ollama server timeout after {}ms", elapsed.as_millis()))
+            }
+            _ = tokio::spawn(async move {
+                loop {
+                    match ollama_api.health().await {
+                        Ok(true) => break,
+                        Ok(false) | Err(_) => tokio::time::sleep(Duration::from_millis(50)).await
+                    }
+                }
+            }) => {
+                Ok(())
             }
-            std::thread::sleep(std::time::Duration::from_millis(500));
-        }
-        if !success {
-            return Err("wait ollama server timeout".to_string());
         }
-        Ok(())
     }
 
     pub async fn spawn(&self, ensure_model: Option<&str>) -> Result<(), String> {

diff --git a/...desktop/src-tauri/src/local_shinkai_node/process_handlers/shinkai_node_process_handler.rs b/...desktop/src-tauri/src/local_shinkai_node/process_handlers/shinkai_node_process_handler.rs
@@ -1,7 +1,7 @@
-use std::fs;
+use std::{fs, time::Duration};
 
 use regex::Regex;
-use tokio::sync::mpsc::Sender;
+use tokio::{sync::mpsc::Sender};
 
 use crate::local_shinkai_node::shinkai_node_options::ShinkaiNodeOptions;
 
@@ -18,7 +18,8 @@ pub struct ShinkaiNodeProcessHandler {
 }
 
 impl ShinkaiNodeProcessHandler {
-    const HEALTH_TIMEOUT_MS: u64 = 500;
+    const HEALTH_REQUEST_TIMEOUT_MS: u64 = 250;
+    const HEALTH_TIMEOUT_MS: u64 = 5000;
     const PROCESS_NAME: &'static str = "shinkai-node";
     const READY_MATCHER: &'static str = "listening on ";
 
@@ -44,12 +45,12 @@ impl ShinkaiNodeProcessHandler {
         base_url
     }
 
-    async fn health(&self) -> Result<bool, ()> {
-        let url = format!("{}/v1/shinkai_health", self.get_base_url());
+    async fn health(base_url: &str, timeout_ms: u64) -> Result<bool, ()> {
+        let url = format!("{}/v1/shinkai_health", base_url);
         let client = reqwest::Client::new();
         if let Ok(response) = client
             .get(&url)
-            .timeout(std::time::Duration::from_millis(400))
+            .timeout(std::time::Duration::from_millis(timeout_ms))
             .send()
             .await
         {
@@ -60,22 +61,25 @@ impl ShinkaiNodeProcessHandler {
     }
 
     async fn wait_shinkai_node_server(&self) -> Result<(), String> {
+        let timeout = Duration::from_millis(Self::HEALTH_TIMEOUT_MS);
         let start_time = std::time::Instant::now();
-        let mut success = false;
-        while std::time::Instant::now().duration_since(start_time)
-            < std::time::Duration::from_millis(Self::HEALTH_TIMEOUT_MS)
-        {
-            let status = self.health().await.unwrap();
-            if status {
-                success = true;
-                break;
+        let base_url = self.get_base_url();
+        tokio::select! {
+            _ = tokio::time::sleep(timeout) => {
+                let elapsed = start_time.elapsed();
+                Err(format!("wait shinkai-node server timeout after {}ms", elapsed.as_millis()))
+            }
+            _ = tokio::spawn(async move {
+                loop {
+                    match Self::health(base_url.as_str(), Self::HEALTH_REQUEST_TIMEOUT_MS).await {
+                        Ok(true) => break,
+                        Ok(false) | Err(_) => tokio::time::sleep(Duration::from_millis(50)).await
+                    }
+                }
+            }) => {
+                Ok(())
             }
-            std::thread::sleep(std::time::Duration::from_millis(500));
-        }
-        if !success {
-            return Err("wait shinkai-node server timeout".to_string());
         }
-        Ok(())
     }
 
     pub fn set_options(&mut self, options: ShinkaiNodeOptions) -> ShinkaiNodeOptions {

diff --git a/apps/shinkai-desktop/src/lib/shinkai-node-manager/shinkai-node-manager-client-types.ts b/apps/shinkai-desktop/src/lib/shinkai-node-manager/shinkai-node-manager-client-types.ts
@@ -93,16 +93,24 @@ export type ShinkaiNodeManagerEventMap =
     };
 
 export type ShinkaiNodeOptions = {
-  port?: number;
-  ws_port?: number;
-  unstructured_server_url?: string;
-  embeddings_server_url?: string;
-  first_device_needs_registration_code?: string;
-  initial_agent_names?: string;
-  initial_agent_urls?: string;
-  initial_agent_models?: string;
-  initial_agent_api_keys?: string;
-  starting_num_qr_devices?: number;
+   node_api_ip?: string,
+   node_api_port?: string,
+   node_ws_port?: string,
+   node_ip?: string,
+   node_port?: string,
+   global_identity_name?: string,
+   node_storage_path?: string,
+   unstructured_server_url?: string,
+   embeddings_server_url?: string,
+   first_device_needs_registration_code?: string,
+   initial_agent_names?: string,
+   initial_agent_urls?: string,
+   initial_agent_models?: string,
+   initial_agent_api_keys?: string,
+   starting_num_qr_devices?: string,
+   log_all?: string,
+   proxy_identity?: string,
+   rpc_url?: string,
 };
 
 export type LogEntry = {

diff --git a/apps/shinkai-desktop/src/store/shinkai-node-manager.ts b/apps/shinkai-desktop/src/store/shinkai-node-manager.ts
@@ -4,7 +4,7 @@ import { devtools, persist } from 'zustand/middleware';
 
 import { ShinkaiNodeOptions } from '../lib/shinkai-node-manager/shinkai-node-manager-client-types';
 import { isLocalShinkaiNode } from '../lib/shinkai-node-manager/shinkai-node-manager-windows-utils';
-import { useAuth } from './auth';
+import { SetupData, useAuth } from './auth';
 
 type ShinkaiNodeManagerStore = {
   isInUse: boolean | null;
@@ -35,12 +35,20 @@ export const useShinkaiNodeManager = create<ShinkaiNodeManagerStore>()(
   ),
 );
 
-useAuth.subscribe((state) => {
-  handleAuthSideEffect(state.auth?.node_address || '');
+useAuth.subscribe((state, prevState) => {
+  handleAuthSideEffect(state.auth, prevState.auth);
 });
 
-const handleAuthSideEffect = async (nodeAddress: string) => {
-  const isLocal = isLocalShinkaiNode(nodeAddress);
-  const isRunning: boolean = await invoke('shinkai_node_is_running');
-  useShinkaiNodeManager.getState().setIsInUse(isLocal && isRunning);
+const handleAuthSideEffect = async (auth: SetupData | null, prevAuth: SetupData | null) => {
+  // SignOut case
+  if (prevAuth && !auth) {
+    useShinkaiNodeManager.getState().setIsInUse(false);
+    return;
+  }
+  // SignIn
+  if (!prevAuth) {
+    const isLocal = isLocalShinkaiNode(auth?.node_address || '');
+    const isRunning: boolean = await invoke('shinkai_node_is_running');
+    useShinkaiNodeManager.getState().setIsInUse(isLocal && isRunning);
+  }
 };
diff --git a/apps/shinkai-desktop/src/windows/shinkai-node-manager/main.tsx b/apps/shinkai-desktop/src/windows/shinkai-node-manager/main.tsx
@@ -203,7 +203,7 @@ const App = () => {
         <img alt="shinkai logo" className="h-10 w-10" src={logo} />
         <div className="ml-4 flex flex-col">
           <span className="text-lg">Local Shinkai Node</span>
-          <span className="text-gray-80 text-sm">{`http://localhost:${shinkaiNodeOptions?.port}`}</span>
+          <span className="text-gray-80 text-sm">{`API URL: http://${shinkaiNodeOptions?.node_api_ip}:${shinkaiNodeOptions?.node_api_port}`}</span>
         </div>
         <div className="flex grow flex-row items-center justify-end space-x-4">
           <TooltipProvider>