feat: added SHINKAI_NODE_LOCATION

dcSpark · Nov 23, 2024 · 014dc5e · 014dc5e
1 parent b445381
commit 014dc5e
Show file tree

Hide file tree

Showing 3 changed files with 111 additions and 14 deletions.
diff --git a/libs/shinkai-tools-runner/src/tools/deno_runner.rs b/libs/shinkai-tools-runner/src/tools/deno_runner.rs
@@ -5,7 +5,10 @@ use tokio::{
 
 use crate::tools::{deno_execution_storage::DenoExecutionStorage, path_buf_ext::PathBufExt};
 
-use super::{container_utils::DockerStatus, deno_runner_options::DenoRunnerOptions};
+use super::{
+    container_utils::DockerStatus,
+    deno_runner_options::{DenoRunnerOptions, RunnerType},
+};
 use std::{
     collections::HashMap,
     path::{self},
@@ -57,12 +60,16 @@ impl DenoRunner {
         envs: Option<HashMap<String, String>>,
         max_execution_timeout: Option<Duration>,
     ) -> anyhow::Result<Vec<String>> {
-        if !self.options.force_deno_in_host
-            && super::container_utils::is_docker_available() == DockerStatus::Running
-        {
-            self.run_in_docker(code, envs, max_execution_timeout).await
-        } else {
-            self.run_in_host(code, envs, max_execution_timeout).await
+        match self.options.force_runner_type {
+            Some(RunnerType::Host) => self.run_in_host(code, envs, max_execution_timeout).await,
+            Some(RunnerType::Docker) => self.run_in_docker(code, envs, max_execution_timeout).await,
+            _ => {
+                if super::container_utils::is_docker_available() == DockerStatus::Running {
+                    self.run_in_docker(code, envs, max_execution_timeout).await
+                } else {
+                    self.run_in_host(code, envs, max_execution_timeout).await
+                }
+            }
         }
     }
 
@@ -133,6 +140,13 @@ impl DenoRunner {
             "DENO_DIR={}",
             execution_storage.relative_to_root(execution_storage.deno_cache.clone())
         ));
+
+        container_envs.push(String::from("-e"));
+        container_envs.push(format!(
+            "SHINKAI_NODE_LOCATION={}://host.docker.internal:{}",
+            self.options.shinkai_node_location.protocol, self.options.shinkai_node_location.port
+        ));
+
         if let Some(envs) = envs {
             for (key, value) in envs {
                 let env = format!("{}={}", key, value);
@@ -293,6 +307,16 @@ impl DenoRunner {
             .kill_on_drop(true);
 
         command.env("DENO_DIR", execution_storage.deno_cache.clone());
+        command.env(
+            "SHINKAI_NODE_LOCATION",
+            format!(
+                "{}://{}:{}",
+                self.options.shinkai_node_location.protocol,
+                self.options.shinkai_node_location.host,
+                self.options.shinkai_node_location.port
+            ),
+        );
+
         if let Some(envs) = envs {
             command.envs(envs);
         }

diff --git a/libs/shinkai-tools-runner/src/tools/deno_runner.test.rs b/libs/shinkai-tools-runner/src/tools/deno_runner.test.rs
@@ -1,6 +1,8 @@
 use crate::tools::{
-    deno_execution_storage::DenoExecutionStorage, deno_runner::DenoRunner,
-    deno_runner_options::DenoRunnerOptions, execution_context::ExecutionContext,
+    deno_execution_storage::DenoExecutionStorage,
+    deno_runner::DenoRunner,
+    deno_runner_options::{DenoRunnerOptions, RunnerType, ShinkaiNodeLocation},
+    execution_context::ExecutionContext,
 };
 use std::collections::HashMap;
 
@@ -55,7 +57,7 @@ async fn test_write_forbidden_folder() {
         .try_init();
 
     let mut deno_runner = DenoRunner::new(DenoRunnerOptions {
-        force_deno_in_host: true,
+        force_runner_type: Some(RunnerType::Host),
         ..Default::default()
     });
 
@@ -113,3 +115,51 @@ async fn test_execution_storage_cache_contains_files() {
     let cache_files = std::fs::read_dir(&storage.deno_cache).unwrap();
     assert!(cache_files.count() > 0);
 }
+
+#[tokio::test]
+async fn test_run_with_shinkai_node_location_host() {
+    let _ = env_logger::builder()
+        .filter_level(log::LevelFilter::Info)
+        .is_test(true)
+        .try_init();
+
+    let mut deno_runner = DenoRunner::new(DenoRunnerOptions {
+        shinkai_node_location: ShinkaiNodeLocation {
+            protocol: String::from("https"),
+            host: String::from("127.0.0.2"),
+            port: 9554,
+        },
+        force_runner_type: Some(RunnerType::Host),
+        ..Default::default()
+    });
+    let code = r#"
+      console.log(process.env.SHINKAI_NODE_LOCATION);
+    "#;
+
+    let result = deno_runner.run(code, None, None).await.unwrap();
+    assert_eq!(result.first().unwrap(), "https://127.0.0.2:9554");
+}
+
+#[tokio::test]
+async fn test_run_with_shinkai_node_location_docker() {
+    let _ = env_logger::builder()
+        .filter_level(log::LevelFilter::Info)
+        .is_test(true)
+        .try_init();
+
+    let mut deno_runner = DenoRunner::new(DenoRunnerOptions {
+        shinkai_node_location: ShinkaiNodeLocation {
+            protocol: String::from("https"),
+            host: String::from("127.0.0.2"),
+            port: 9554,
+        },
+        force_runner_type: Some(RunnerType::Docker),
+        ..Default::default()
+    });
+    let code = r#"
+      console.log(process.env.SHINKAI_NODE_LOCATION);
+    "#;
+
+    let result = deno_runner.run(code, None, None).await.unwrap();
+    assert_eq!(result.first().unwrap(), "https://host.docker.internal:9554");
+}
diff --git a/libs/shinkai-tools-runner/src/tools/deno_runner_options.rs b/libs/shinkai-tools-runner/src/tools/deno_runner_options.rs
@@ -2,12 +2,26 @@ use std::path::PathBuf;
 
 use super::execution_context::ExecutionContext;
 
+#[derive(Clone)]
+pub struct ShinkaiNodeLocation {
+    pub protocol: String,
+    pub host: String,
+    pub port: u16,
+}
+
+#[derive(Clone)]
+pub enum RunnerType {
+    Host,
+    Docker,
+}
+
 #[derive(Clone)]
 pub struct DenoRunnerOptions {
     pub context: ExecutionContext,
     pub deno_binary_path: PathBuf,
     pub deno_image_name: String,
-    pub force_deno_in_host: bool,
+    pub force_runner_type: Option<RunnerType>,
+    pub shinkai_node_location: ShinkaiNodeLocation,
 }
 
 impl Default for DenoRunnerOptions {
@@ -20,9 +34,18 @@ impl Default for DenoRunnerOptions {
             } else {
                 "./shinkai-tools-runner-resources/deno"
             }),
-            force_deno_in_host: std::env::var("CI_FORCE_DENO_IN_HOST")
-                .map(|val| val == "true")
-                .unwrap_or(false),
+            force_runner_type: std::env::var("CI_FORCE_RUNNER_TYPE")
+                .map(|val| match val.as_str() {
+                    "host" => Some(RunnerType::Host),
+                    "docker" => Some(RunnerType::Docker),
+                    _ => None,
+                })
+                .unwrap_or(None),
+            shinkai_node_location: ShinkaiNodeLocation {
+                protocol: String::from("http"),
+                host: String::from("127.0.0.1"),
+                port: 9550,
+            },
         }
     }
 }