Update tasks

1
package main
2
3
import (
4
	"context"
5
	"fmt"
6
	"net/http"
7
	"os"
8
	"time"
9
10
	Lattice "github.com/anduril/lattice-sdk-go/v2"
11
	"github.com/anduril/lattice-sdk-go/v2/client"
12
	"github.com/anduril/lattice-sdk-go/v2/option"
13
)
14
15
func main() {
16
	// Get environment variables
17
	latticeEndpoint := os.Getenv("LATTICE_ENDPOINT")
18
	environmentToken := os.Getenv("ENVIRONMENT_TOKEN")
19
	sandboxesToken := os.Getenv("SANDBOXES_TOKEN")
20
21
	// Check required environment variables
22
	if latticeEndpoint == "" || environmentToken == "" || sandboxesToken == "" {
23
		fmt.Println("Missing required environment variables")
24
		os.Exit(1)
25
	}
26
27
	// Initialize headers for sandbox authorization
28
	headers := http.Header{}
29
	headers.Add("Anduril-Sandbox-Authorization", fmt.Sprintf("Bearer %s", sandboxesToken))
30
	// Create the client
31
	LatticeClient := client.NewClient(
32
		option.WithToken(environmentToken),
33
		option.WithBaseURL(fmt.Sprintf("https://%s", latticeEndpoint)),
34
		option.WithHTTPHeader(headers),
35
	)
36
37
	// Set the entity ID to listen for tasks
38
	entityId := "<ENTITY_ID>"
39
	fmt.Printf("Listening for tasks for entity %s...\n", entityId)
40
41
	// Create context for the request
42
	ctx := context.Background()
43
44
	// Main loop to continuously listen for tasks
45
	for {
46
		try := func() bool {
47
			// Listen for tasks using the listenAsAgent API
48
			agentSelector := Lattice.AgentListener{
49
				AgentSelector: &Lattice.EntityIDsSelector{
50
					EntityIDs: []string{entityId},
51
				},
52
			}
53
			response, err := LatticeClient.Tasks.ListenAsAgent(ctx, &agentSelector)
54
55
			// Handle errors
56
			if err != nil {
57
				fmt.Printf("Error listening for tasks: %v\n", err)
58
				return false
59
			}
60
61
			// Process the request based on its type
62
			if response.ExecuteRequest != nil && response.ExecuteRequest.Task != nil &&
63
				response.ExecuteRequest.Task.Version != nil && response.ExecuteRequest.Task.Version.TaskID != nil {
64
				// Extract task-specific information
65
				taskId := *response.ExecuteRequest.Task.Version.TaskID
66
				taskStatusVersion := *response.ExecuteRequest.Task.Version.StatusVersion
67
68
				fmt.Printf("Starting task %s, version %d\n", taskId, taskStatusVersion)
69
70
				// Update task status to STATUS_EXECUTING
71
				result, err := startTask(ctx, LatticeClient, taskId, taskStatusVersion, entityId)
72
				if err != nil {
73
					fmt.Printf("Error starting task: %v\n", err)
74
					return false
75
				}
76
77
				fmt.Printf("Started task with status version: %d\n", *result.StatusVersion)
78
79
			} else if response.CompleteRequest != nil && response.CompleteRequest.TaskID != nil {
80
				fmt.Printf("Completing task: %s\n", *response.CompleteRequest.TaskID)
81
			} else if response.CancelRequest != nil && response.CancelRequest.TaskID != nil {
82
				fmt.Printf("Cancelling task: %s\n", *response.CancelRequest.TaskID)
83
			}
84
85
			return true
86
		}()
87
88
		if !try {
89
			// If there was an error, wait a bit before retrying
90
			time.Sleep(1 * time.Second)
91
		}
92
	}
93
}
94
95
// startTask updates the task status to STATUS_EXECUTING
96
func startTask(ctx context.Context, client *client.Client, taskId string, taskStatusVersion int, agentEntityId string) (*Lattice.TaskVersion, error) {
97
	// Increment status version for the update
98
	taskStatusVersion++
99
100
	// Create system principal with the agent entity ID
101
	principal := Lattice.Principal{
102
		System: &Lattice.System{
103
			EntityID: &agentEntityId,
104
		},
105
	}
106
107
	taskStatus := Lattice.TaskStatus{
108
		Status: Lattice.TaskStatusStatusStatusExecuting.Ptr(),
109
	}
110
111
	// Create task status update request
112
	taskStatusUpdate := Lattice.TaskStatusUpdate{
113
		StatusVersion: &taskStatusVersion,
114
		NewStatus:     &taskStatus,
115
		Author:        &principal,
116
	}
117
118
	// Call the UpdateTaskStatus API
119
	task, err := client.Tasks.UpdateTaskStatus(ctx, taskId, &taskStatusUpdate)
120
	if err != nil {
121
		return nil, fmt.Errorf("error updating task status: %w", err)
122
	}
123
124
	return task.Version, nil
125
}