Update tasks

1
package main
2
3
import (
4
	"context"
5
	"errors"
6
	"fmt"
7
	"io"
8
	"log"
9
	"net/http"
10
	"os"
11
	"time"
12
13
	Lattice "github.com/anduril/lattice-sdk-go/v4"
14
	"github.com/anduril/lattice-sdk-go/v4/client"
15
	"github.com/anduril/lattice-sdk-go/v4/option"
16
)
17
18
func main() {
19
	// Get environment variables
20
	latticeEndpoint := os.Getenv("LATTICE_ENDPOINT")
21
	clientSecret := os.Getenv("LATTICE_CLIENT_SECRET")
22
	clientId := os.Getenv("LATTICE_CLIENT_ID")
23
	sandboxesToken := os.Getenv("SANDBOXES_TOKEN")
24
25
	// Check required environment variables
26
	if latticeEndpoint == "" || clientId == "" || clientSecret == "" || sandboxesToken == "" {
27
		fmt.Println("Missing required environment variables")
28
		os.Exit(1)
29
	}
30
31
	// Initialize headers for sandbox authorization
32
	headers := http.Header{}
33
	headers.Add("Anduril-Sandbox-Authorization", fmt.Sprintf("Bearer %s", sandboxesToken))
34
	// Create the client
35
	LatticeClient := client.NewClient(
36
		option.WithClientCredentials(clientId, clientSecret),
37
		option.WithBaseURL(fmt.Sprintf("https://%s", latticeEndpoint)),
38
		option.WithHTTPHeader(headers),
39
	)
40
41
	// Set the entity ID to listen for tasks
42
	entityId := "<AGENT_ID>"
43
	fmt.Printf("Streaming tasks for entity %s...\n", entityId)
44
45
	// Create context for the request
46
	ctx := context.Background()
47
48
	// Create agent stream request
49
	agentStreamRequest := Lattice.AgentStreamRequest{
50
		AgentSelector: &Lattice.EntityIDsSelector{
51
			EntityIDs: []string{entityId},
52
		},
53
	}
54
55
	// Stream tasks
56
	stream, err := LatticeClient.Tasks.StreamAsAgent(ctx, &agentStreamRequest)
57
	if err != nil {
58
		fmt.Printf("Error streaming tasks: %v\n", err)
59
		os.Exit(1)
60
	}
61
62
	// Process stream events
63
	for {
64
		select {
65
		case <-ctx.Done():
66
			log.Printf("Context canceled: %v", ctx.Err())
67
			return
68
		default:
69
			// Continue processing
70
		}
71
72
		event, err := stream.Recv()
73
74
		if errors.Is(err, io.EOF) {
75
			log.Println("Stream completed successfully.")
76
			return
77
		}
78
79
		if err != nil {
80
			log.Printf("Error receiving message: %v", err)
81
			continue
82
		}
83
84
		if event.Event == "heartbeat" {
85
			timestamp := *event.Heartbeat.Timestamp
86
			log.Printf("Heartbeat: %s", timestamp)
87
		} else {
88
			request := event.GetAgentRequest()
89
			if executeRequest := request.GetExecuteRequest(); executeRequest != nil {
90
				task := executeRequest.GetTask()
91
				if task != nil {
92
					taskId := *task.GetVersion().GetTaskID()
93
					taskStatusVersion := *task.GetVersion().GetStatusVersion()
94
					description := *task.GetDescription()
95
96
					log.Printf("Starting task %s, version %d: %s", taskId, taskStatusVersion, description)
97
98
					// Update task status to STATUS_EXECUTING
99
					result, err := startTask(ctx, LatticeClient, taskId, int(taskStatusVersion), entityId)
100
					if err != nil {
101
						log.Printf("Error starting task: %v", err)
102
						continue
103
					}
104
105
					log.Printf("Started task with status version: %d", *result.StatusVersion)
106
				}
107
			} else if completeRequest := request.GetCompleteRequest(); completeRequest != nil {
108
				taskToComplete := completeRequest.GetTaskID()
109
				if taskToComplete != nil {
110
					log.Printf("Completing task: %s", *taskToComplete)
111
				}
112
			} else if cancelRequest := request.GetCancelRequest(); cancelRequest != nil {
113
				taskToCancel := cancelRequest.GetTaskID()
114
				if taskToCancel != nil {
115
					log.Printf("Cancelling task: %s", *taskToCancel)
116
				}
117
			}
118
		}
119
120
		// Sleep briefly to prevent tight looping
121
		time.Sleep(100 * time.Millisecond)
122
	}
123
}
124
125
// startTask updates the task status to STATUS_EXECUTING
126
func startTask(ctx context.Context, client *client.Client, taskId string, taskStatusVersion int, agentEntityId string) (*Lattice.TaskVersion, error) {
127
	// Increment status version for the update
128
	taskStatusVersion++
129
130
	// Create system principal with the agent entity ID
131
	principal := Lattice.Principal{
132
		System: &Lattice.System{
133
			EntityID: &agentEntityId,
134
		},
135
	}
136
137
	taskStatus := Lattice.TaskStatus{
138
		Status: Lattice.TaskStatusStatusStatusExecuting.Ptr(),
139
	}
140
141
	// Create task status update request
142
	taskStatusUpdate := Lattice.TaskStatusUpdate{
143
		TaskID:        taskId,
144
		StatusVersion: &taskStatusVersion,
145
		NewStatus:     &taskStatus,
146
		Author:        &principal,
147
	}
148
149
	// Call the UpdateTaskStatus API
150
	task, err := client.Tasks.UpdateTaskStatus(ctx, &taskStatusUpdate)
151
	if err != nil {
152
		return nil, fmt.Errorf("error updating task status: %w", err)
153
	}
154
155
	return task.Version, nil
156
}