From 88af5b838ca063564ed7aabf520c298f6e9f9a8f Mon Sep 17 00:00:00 2001
From: peefy <xpf6677@163.com>
Date: Mon, 19 Feb 2024 10:27:27 +0800
Subject: [PATCH] feat: enhance plugin agent flag in APIs

Signed-off-by: peefy <xpf6677@163.com>
---
 kclvm/api/src/service/service_impl.rs | 23 +++++++++--------------
 kclvm/api/src/service/util.rs         |  7 +++++--
 2 files changed, 14 insertions(+), 16 deletions(-)

diff --git a/kclvm/api/src/service/service_impl.rs b/kclvm/api/src/service/service_impl.rs
index 5b539aa80..18c37a862 100644
--- a/kclvm/api/src/service/service_impl.rs
+++ b/kclvm/api/src/service/service_impl.rs
@@ -286,25 +286,20 @@ impl KclvmServiceImpl {
     ///     ..Default::default()
     /// };
     /// let error = serv.exec_program(args).unwrap_err();
-    /// assert!(error.contains("Cannot find the kcl file"), "{error}");
+    /// assert!(error.to_string().contains("Cannot find the kcl file"), "{error}");
     ///
     /// let args = &ExecProgramArgs {
     ///     k_filename_list: vec![],
     ///     ..Default::default()
     /// };
     /// let error = serv.exec_program(args).unwrap_err();
-    /// assert!(error.contains("No input KCL files or paths"), "{error}");
+    /// assert!(error.to_string().contains("No input KCL files or paths"), "{error}");
     /// ```
-    pub fn exec_program(&self, args: &ExecProgramArgs) -> Result<ExecProgramResult, String> {
+    pub fn exec_program(&self, args: &ExecProgramArgs) -> anyhow::Result<ExecProgramResult> {
         // transform args to json
-        let args_json = serde_json::to_string(args).unwrap();
-
+        let exec_args = transform_exec_para(&Some(args.clone()), self.plugin_agent)?;
         let sess = ParseSessionRef::default();
-        let result = exec_program(
-            sess,
-            &kclvm_runner::ExecProgramArgs::from_str(args_json.as_str()),
-        )
-        .map_err(|err| err.to_string())?;
+        let result = exec_program(sess, &exec_args)?;
 
         Ok(ExecProgramResult {
             json_result: result.json_result,
@@ -336,7 +331,7 @@ impl KclvmServiceImpl {
     /// assert!(!artifact.path.is_empty());
     /// ```
     pub fn build_program(&self, args: &BuildProgramArgs) -> anyhow::Result<BuildProgramResult> {
-        let exec_args = transform_exec_para(&args.exec_args)?;
+        let exec_args = transform_exec_para(&args.exec_args, self.plugin_agent)?;
         let artifact = build_program(
             ParseSessionRef::default(),
             &exec_args,
@@ -375,7 +370,7 @@ impl KclvmServiceImpl {
     /// assert_eq!(exec_result.yaml_result, "alice:\n  age: 18");
     /// ```
     pub fn exec_artifact(&self, args: &ExecArtifactArgs) -> anyhow::Result<ExecProgramResult> {
-        let exec_args = transform_exec_para(&args.exec_args)?;
+        let exec_args = transform_exec_para(&args.exec_args, self.plugin_agent)?;
         let result = exec_artifact(&args.path, &exec_args)?;
         Ok(ExecProgramResult {
             json_result: result.json_result,
@@ -507,7 +502,7 @@ impl KclvmServiceImpl {
         args: &GetFullSchemaTypeArgs,
     ) -> anyhow::Result<GetSchemaTypeResult> {
         let mut type_list = Vec::new();
-        let exec_args = transform_exec_para(&args.exec_args)?;
+        let exec_args = transform_exec_para(&args.exec_args, self.plugin_agent)?;
         for (_k, schema_ty) in get_full_schema_type(
             Some(&args.schema_name),
             CompilationOptions {
@@ -893,7 +888,7 @@ impl KclvmServiceImpl {
     /// ```
     pub fn test(&self, args: &TestArgs) -> anyhow::Result<TestResult> {
         let mut result = TestResult::default();
-        let exec_args = transform_exec_para(&args.exec_args)?;
+        let exec_args = transform_exec_para(&args.exec_args, self.plugin_agent)?;
         let opts = testing::TestOptions {
             exec_args,
             run_regexp: args.run_regexp.clone(),
diff --git a/kclvm/api/src/service/util.rs b/kclvm/api/src/service/util.rs
index 2d4247da2..f74b5fa1f 100644
--- a/kclvm/api/src/service/util.rs
+++ b/kclvm/api/src/service/util.rs
@@ -13,12 +13,15 @@ pub(crate) fn transform_str_para(para: &str) -> Option<String> {
 #[inline]
 pub(crate) fn transform_exec_para(
     exec_args: &Option<ExecProgramArgs>,
+    plugin_agent: u64,
 ) -> anyhow::Result<kclvm_runner::ExecProgramArgs> {
-    Ok(match exec_args {
+    let mut args = match exec_args {
         Some(exec_args) => {
             let args_json = serde_json::to_string(exec_args)?;
             kclvm_runner::ExecProgramArgs::from_str(args_json.as_str())
         }
         None => kclvm_runner::ExecProgramArgs::default(),
-    })
+    };
+    args.plugin_agent = plugin_agent;
+    Ok(args)
 }