Added prompt and app tool mapping from getTemplateContext() for agents

Author: frikky

ai.go

@@ -7654,21 +7654,23 @@ func ReduceAgentResponseData(rawResponse []byte, dataFilter string, fieldsNeeded
 
 // createNextActions = false => start of agent to find initial decisions
 // createNextActions = true => mid-agent to decide next steps
-func getTemplateContext(ctx context.Context, template string, execution WorkflowExecution) (string, string, error) {
+func getTemplateContext(ctx context.Context, template string, execution WorkflowExecution) (string, string, []string, error) {
+	// FIXME: Handle dynamic templates here too, based on user input
+
 	switch template {
 	case "workflow-edit":
 		return buildWorkflowEditContext(ctx, execution)
 	case "computer-use":
 		return buildComputerUseContext(ctx, execution)
 	default:
-		return "", "", nil
+		return "", "", []string{}, nil
 	}
 }
 
-func buildComputerUseContext(ctx context.Context, execution WorkflowExecution) (string, string, error) {
+func buildComputerUseContext(ctx context.Context, execution WorkflowExecution) (string, string, []string, error) {
 	systemRule := `You are a computer use agent that can execute commands to control a computer. Your goal is to use the Terminal, Keyboard, Mouse and Screenshots to perform the task the user intends in the best possible way. Make assumptions for what they most likely want to perform, and continue it is done.
 
-Use the 'post_control_mouse_and_keyboard' function for keyboard & mouse control if it is available. You can chain together escaped JSON commands in the the "actions" array using the operations detailed below.
+Use the 'post_control_mouse_and_keyboard' function for keyboard & mouse control if it is available. You can chain together escaped JSON commands in the the "actions" array using the operations detailed below. If an action takes more than 30 seconds, it will return an execution_id and authorization key to be used for polling results. When polling, always add a 30 second delay. 
 
 Params for each keyboard & mouse operation: 
 1. keyboard.press: {\"op\":\"keyboard.press\",\"params\":{\"key\":75}}
@@ -7691,13 +7693,18 @@ When sending body to the post_control_mouse_and_keyboard function, the following
 }
 '''
 
-Prioritise terminal commands before screenshots. IF an action takes more than 30 seconds, it will return an execution_id and authorization key to be used for polling results. ALWAYS validate if the action was successful by checking the output of the or a before and after screenshot. If it was not successful, you must try again with a different approach.
+ALWAYS validate if the action was successful by checking the output of the operations before AND after with screenshots or terminal input/output. If it was not successful, you must try again with a different approach.
 	`
 
-	return systemRule, "", nil
+	templateContext := ""
+	requiredApps := []string{
+		"app:48a954b9440b3913b8a2620e57b94a75:shuffle_host_monitors",
+	}
+
+	return systemRule, templateContext, requiredApps, nil
 }
 
-func buildWorkflowEditContext(ctx context.Context, execution WorkflowExecution) (string, string, error) {
+func buildWorkflowEditContext(ctx context.Context, execution WorkflowExecution) (string, string, []string, error) {
 	targetWorkflowId := execution.ExecutionArgument
 
 	user := User{
@@ -7711,7 +7718,7 @@ func buildWorkflowEditContext(ctx context.Context, execution WorkflowExecution)
 	appsJson, err := json.Marshal(appSummaries)
 	if err != nil {
 		log.Printf("[WARNING] buildWorkflowEditContext: failed marshaling app summaries for org %s: %s", execution.ExecutionOrg, err)
-		return "", "", fmt.Errorf("failed to marshal app summaries: %w", err)
+		return "", "", []string{}, fmt.Errorf("failed to marshal app summaries: %w", err)
 	}
 
 	// What we tell the agent about its workflow_id
@@ -7910,7 +7917,12 @@ CRITICAL RULES FOR THE AGENT
 %s
 <End of Available Apps>`, workflowIdLine, string(appsJson))
 
-	return systemRule, templateContext, nil
+	requiredApps := []string{
+		"app:7db43ccd25261967b095cfbd467a75cc:shuffle_apps",
+		"app:de4ef2287bd41b9d5563e39989643ee6:shuffle_workflows_builder",
+	}
+
+	return systemRule, templateContext, requiredApps, nil
 }
 
 func getWorkflowEditPromptRemovals() []string {
@@ -8856,8 +8868,9 @@ data_filter:
 	// If a template is set, get secondary system rules + extra context.
 	templateSystemRule := ""
 	templateContext := ""
+	requiredApps := []string{}
 	if len(template) > 0 {
-		templateSystemRule, templateContext, err = getTemplateContext(ctx, template, execution)
+		templateSystemRule, templateContext, requiredApps, err = getTemplateContext(ctx, template, execution)
 		if err != nil {
 			log.Printf("[ERROR] Failed to get template context: %v", err)
 		}
@@ -8875,6 +8888,18 @@ data_filter:
 		agentReasoningEffort = foundReasoning
 	}
 
+	added := false
+	for _, requiredApp := range requiredApps { 
+		if !strings.Contains(allowedActionString, requiredApp) {
+			added = true
+			allowedActionString += "," + requiredApp 
+		}
+	}
+
+	if added {
+		allowedActionString = strings.TrimPrefix(allowedActionString, ",")
+	}
+
 	agentOutput := AgentOutput{
 		Status:         "RUNNING",
 		Input:          userMessage,
@@ -8950,7 +8975,7 @@ data_filter:
 	if len(templateContext) > 0 {
 		completionRequest.Messages = append(completionRequest.Messages, openai.ChatCompletionMessage{
 			//Role:    openai.ChatMessageRoleUser,
-			Role:    openai.ChatMessageRoleSystem,
+			Role:    openai.ChatMessageRoleUser,
 			Content: "SKILL CONTEXT:\n" + templateContext,
 		})
 	}

db-connector.go

@@ -2446,6 +2446,19 @@ func GetEnvironment(ctx context.Context, id, orgId string) (*Environment, error)
 			cacheData := []byte(cache.([]uint8))
 			err = json.Unmarshal(cacheData, &env)
 			if err == nil {
+
+				timenow := time.Now().Unix()
+				if env.SensorGroup { 
+					for sensorIndex, _ := range env.SensorHosts {
+						sensor := env.SensorHosts[sensorIndex]
+							
+						env.SensorHosts[sensorIndex].Active = false 
+						if sensor.Checkin > 0 && timenow-sensor.Checkin < 300 {
+							env.SensorHosts[sensorIndex].Active = true
+						}
+					}
+				}
+
 				return env, nil
 			}
 		} else {
@@ -2566,16 +2579,26 @@ func GetEnvironment(ctx context.Context, id, orgId string) (*Environment, error)
 	} else {
 		key := datastore.NameKey(nameKey, strings.ToLower(id), nil)
 		if err := project.Dbclient.Get(ctx, key, env); err != nil {
+			log.Printf("[ERROR] Problem in environment loading of %s", id)
 			if strings.Contains(err.Error(), `cannot load field`) {
-				log.Printf("[INFO] Error in environment loading of %s", id)
 				err = nil
 			} else {
 				return env, err
 			}
 		}
 	}
 
-	//log.Printf("[DEBUG] Got hit: %s", env)
+	timenow := time.Now().Unix()
+	if env.SensorGroup { 
+		for sensorIndex, _ := range env.SensorHosts {
+			sensor := env.SensorHosts[sensorIndex]
+				
+			env.SensorHosts[sensorIndex].Active = false 
+			if sensor.Checkin > 0 && timenow-sensor.Checkin < 300 {
+				env.SensorHosts[sensorIndex].Active = true
+			}
+		}
+	}
 
 	if project.CacheDb {
 		//log.Printf("[DEBUG] Setting cache for workflow %s", cacheKey)
@@ -6414,6 +6437,20 @@ func GetEnvironments(ctx context.Context, orgId string) ([]Environment, error) {
 				//	log.Printf("[DEBUG] Got %d environments from cache for orgId '%s'", len(environments), orgId)
 				//}
 
+				timenow := time.Now().Unix()
+				for envIndex, env := range environments {
+					if env.SensorGroup { 
+						for sensorIndex, _ := range env.SensorHosts {
+							sensor := env.SensorHosts[sensorIndex]
+								
+							environments[envIndex].SensorHosts[sensorIndex].Active = false 
+							if sensor.Checkin > 0 && timenow-sensor.Checkin < 300 {
+								environments[envIndex].SensorHosts[sensorIndex].Active = true
+							}
+						}
+					}
+				}
+
 				return environments, nil
 			}
 		} else {
@@ -6584,6 +6621,18 @@ func GetEnvironments(ctx context.Context, orgId string) ([]Environment, error) {
 	// Fixing environment return search problems
 	timenow := time.Now().Unix()
 	for envIndex, env := range environments {
+
+		if env.SensorGroup { 
+			for sensorIndex, _ := range env.SensorHosts {
+				sensor := env.SensorHosts[sensorIndex]
+					
+				environments[envIndex].SensorHosts[sensorIndex].Active = false 
+				if sensor.Checkin > 0 && timenow-sensor.Checkin < 300 {
+					environments[envIndex].SensorHosts[sensorIndex].Active = true
+				}
+			}
+		}
+
 		if env.Name == "Cloud" {
 			environments[envIndex].Type = "cloud"
 			environments[envIndex].RunType = "cloud"

executions.go

@@ -76,6 +76,14 @@ func Fixexecution(ctx context.Context, workflowExecution WorkflowExecution) (Wor
 
 			// Special cleanup for agents
 			if innerresult.Action.AppName == "AI Agent" || innerresult.Action.AppName == "Shuffle Agent" {
+				if workflowExecution.Status == "FINISHED" || workflowExecution.Status == "ABORTED" { 
+					if workflowExecution.Status == "FINISHED" {
+						log.Printf("[DEBUG][%s] Fixexecution: Agent execution is finished, skipping agent result %s", workflowExecution.ExecutionId, innerresult.Action.ID)
+					}
+
+					break
+				}
+
 				actionCacheId := fmt.Sprintf("%s_%s_result", workflowExecution.ExecutionId, innerresult.Action.ID)
 				if cachedData, cacheErr := GetCache(ctx, actionCacheId); cacheErr == nil {
 					cachedBytes := []byte(cachedData.([]uint8))

shared.go

@@ -36767,139 +36767,6 @@ func GetWorkflowMinimal(resp http.ResponseWriter, request *http.Request) {
 	resp.Write(responseData)
 }
 
-func AgentWorkflowEditor(resp http.ResponseWriter, request *http.Request) {
-	cors := HandleCors(resp, request)
-	if cors {
-		return
-	}
-
-	ctx := GetContext(request)
-	user, userErr := HandleApiAuthentication(resp, request)
-	if userErr != nil {
-		log.Printf("[AUDIT] Api authentication failed in AgentWorkflowEditor: %s", userErr)
-		resp.WriteHeader(401)
-		resp.Write([]byte(`{"success": false}`))
-		return
-	}
-
-	body, err := ioutil.ReadAll(request.Body)
-	if err != nil {
-		log.Printf("[WARNING] Failed reading body in AgentWorkflowEditor: %s", err)
-		resp.WriteHeader(400)
-		resp.Write([]byte(`{"success": false}`))
-		return
-	}
-
-	var req MCPRequest 
-	err = json.Unmarshal(body, &req)
-
-	if err != nil || len(strings.TrimSpace(req.Params.Input.Text)) == 0 {
-		log.Printf("[WARNING] Bad body in AgentWorkflowEditor. Error: %v, Body: %s", err, string(body))
-		resp.WriteHeader(400)
-		resp.Write([]byte(`{"success": false, "reason": "input field is required"}`))
-		return
-	}
-
-
-	// All context building (workflow state, app actions, rules) is handled inside
-	// buildWorkflowEditContext which is called by getTemplateContext inside HandleAiAgentExecutionStart
-	toolApps := "app:7db43ccd25261967b095cfbd467a75cc:shuffle_apps,app:de4ef2287bd41b9d5563e39989643ee6:shuffle_workflows_builder"
-
-	action := Action{
-		ID:             uuid.NewV4().String(),
-		Name:           "agent",
-		AppName:        "AI Agent",
-		AppID:          "shuffle_agent",
-		AppVersion:     "1.0.0",
-		Environment:    "cloud",
-		Parameters: []WorkflowAppActionParameter{
-			{
-				Name:  "input",
-				Value: req.Params.Input.Text,
-			},
-			{
-				Name:  "action",
-				Value: toolApps,
-			},
-			{
-				Name:  "template",
-				Value: "workflow-edit",
-			},
-			{
-				Name: "execution_mode",
-				Value: "direct",
-			},
-		},
-	}
-
-	workflowId := uuid.NewV4().String()
-	action.SourceWorkflow = workflowId
-
-	exec := WorkflowExecution{
-		Workflow: Workflow{
-			ID: workflowId,
-			Actions: []Action{
-				action,
-			},
-			OrgId:     user.ActiveOrg.Id,
-			Owner:     user.Username,
-			UpdatedBy: user.Username,
-			Start:     action.ID,
-		},
-		Type:             "AGENT",
-		Start:            action.ID,
-		Status:           "EXECUTING",
-		WorkflowId:       workflowId,
-		ExecutionId:      workflowId,
-		ExecutionOrg:     user.ActiveOrg.Id,
-		StartedAt:        int64(time.Now().Unix()),
-		Authorization:    uuid.NewV4().String(),
-		// Store the target workflow_id here so buildWorkflowEditContext can fetch it
-		ExecutionArgument: req.Params.Input.WorkflowId,
-	}
-
-	SetWorkflowExecution(ctx, exec, true)
-
-	log.Printf("[INFO] AgentWorkflowEditor: calling HandleAiAgentExecutionStart for user %s (%s), target_workflow_id=%s, execution_id=%s", user.Username, user.Id, req.Params.Input.WorkflowId, exec.ExecutionId)
-
-	returnAction, err := HandleAiAgentExecutionStart(exec, action, false, "AgentWorkflowEditor")
-	if err != nil {
-		log.Printf("[ERROR] HandleAiAgentExecutionStart failed in AgentWorkflowEditor: %s", err)
-		resp.WriteHeader(500)
-		resp.Write([]byte(fmt.Sprintf(`{"success": false, "reason": "%s"}`, err.Error())))
-		return
-	}
-
-	// Fetch the updated execution to return
-	newExec, err := GetWorkflowExecution(ctx, exec.ExecutionId)
-	if err != nil {
-		log.Printf("[ERROR] Failed to get workflow execution after agent start in AgentWorkflowEditor: %s", err)
-		resp.WriteHeader(500)
-		resp.Write([]byte(`{"success": false}`))
-		return
-	}
-
-	_ = returnAction
-
-	respObj := agentResponse{
-		Success:       true,
-		ExecutionId:   newExec.ExecutionId,
-		Authorization: newExec.Authorization,
-	}
-
-	responseData, err := json.Marshal(respObj)
-	if err != nil {
-		log.Printf("[ERROR] Failed marshalling execution response in AgentWorkflowEditor: %s", err)
-		resp.WriteHeader(500)
-		resp.Write([]byte(`{"success": false}`))
-		return
-	}
-
-	resp.Header().Set("Content-Type", "application/json")
-	resp.WriteHeader(200)
-	resp.Write(responseData)
-}
-
 func generateNodeID() string {
 	return uuid.NewV4().String()
 }

structs.go

@@ -5447,6 +5447,7 @@ type Parsed struct {
 type SensorDetails struct {
 	SensorMode bool   `json:"sensor_mode,omitempty" datastore:"sensor_mode"`
 	Checkin    int64  `json:"checkin" datastore:"checkin"`
+	Active     bool   `json:"active,omitempty" datastore:"active"`
 	Uuid       string `json:"uuid" datastore:"uuid"`
 
 	User     string `json:"user,omitempty" datastore:"user"`
@@ -5470,7 +5471,8 @@ type SensorDetails struct {
 
 // Related to Orborus Agent Mode. Used locally.
 type SensorMode struct {
-	Enabled bool `json:"enabled" datastore:"enabled"`
+	Enabled  bool   `json:"enabled" datastore:"enabled"`
+	Hostname string `json:"hostname" datastore:"hostname"`
 
 	// Compliance
 	ProcessListEnabled  string `json:"process_list_enabled" datastore:"process_list_enabled"`