feat(agentic): type fleet compute budgets
Co-Authored-By: Virgil <virgil@lethean.io>
This commit is contained in:
Virgil
parent
02aea97b7d
commit
8712c7c921
3 changed files with 181 additions and 4 deletions
|
|
@ -11,6 +11,16 @@ import (
|
|||
core "dappco.re/go/core"
|
||||
)
|
||||
|
||||
// node := agentic.FleetNode{AgentID: "charon", Platform: "linux", Status: "online"}
|
||||
type ComputeBudget struct {
|
||||
MaxDailyHours float64 `json:"max_daily_hours,omitempty"`
|
||||
MaxWeeklyCostUSD float64 `json:"max_weekly_cost_usd,omitempty"`
|
||||
QuietStart string `json:"quiet_start,omitempty"`
|
||||
QuietEnd string `json:"quiet_end,omitempty"`
|
||||
PreferModels []string `json:"prefer_models,omitempty"`
|
||||
AvoidModels []string `json:"avoid_models,omitempty"`
|
||||
}
|
||||
|
||||
// node := agentic.FleetNode{AgentID: "charon", Platform: "linux", Status: "online"}
|
||||
type FleetNode struct {
|
||||
ID int `json:"id"`
|
||||
|
|
@ -20,7 +30,7 @@ type FleetNode struct {
|
|||
Models []string `json:"models,omitempty"`
|
||||
Capabilities []string `json:"capabilities,omitempty"`
|
||||
Status string `json:"status"`
|
||||
ComputeBudget map[string]any `json:"compute_budget,omitempty"`
|
||||
ComputeBudget *ComputeBudget `json:"compute_budget,omitempty"`
|
||||
CurrentTaskID int `json:"current_task_id,omitempty"`
|
||||
LastHeartbeatAt string `json:"last_heartbeat_at,omitempty"`
|
||||
RegisteredAt string `json:"registered_at,omitempty"`
|
||||
|
|
@ -710,13 +720,77 @@ func parseFleetNode(values map[string]any) FleetNode {
|
|||
Models: listValue(values["models"]),
|
||||
Capabilities: listValue(values["capabilities"]),
|
||||
Status: stringValue(values["status"]),
|
||||
ComputeBudget: anyMapValue(values["compute_budget"]),
|
||||
ComputeBudget: computeBudgetFromValue(values["compute_budget"]),
|
||||
CurrentTaskID: intValue(values["current_task_id"]),
|
||||
LastHeartbeatAt: stringValue(values["last_heartbeat_at"]),
|
||||
RegisteredAt: stringValue(values["registered_at"]),
|
||||
}
|
||||
}
|
||||
|
||||
func computeBudgetFromValue(value any) *ComputeBudget {
|
||||
switch typed := value.(type) {
|
||||
case *ComputeBudget:
|
||||
if typed == nil || computeBudgetIsZero(*typed) {
|
||||
return nil
|
||||
}
|
||||
return typed
|
||||
case ComputeBudget:
|
||||
if computeBudgetIsZero(typed) {
|
||||
return nil
|
||||
}
|
||||
return &typed
|
||||
case map[string]any:
|
||||
return computeBudgetFromMap(typed)
|
||||
case map[string]string:
|
||||
values := make(map[string]any, len(typed))
|
||||
for key, item := range typed {
|
||||
values[key] = item
|
||||
}
|
||||
return computeBudgetFromMap(values)
|
||||
case string:
|
||||
trimmed := core.Trim(typed)
|
||||
if trimmed == "" {
|
||||
return nil
|
||||
}
|
||||
if core.HasPrefix(trimmed, "{") {
|
||||
var values map[string]any
|
||||
if result := core.JSONUnmarshalString(trimmed, &values); result.OK {
|
||||
return computeBudgetFromMap(values)
|
||||
}
|
||||
}
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
func computeBudgetFromMap(values map[string]any) *ComputeBudget {
|
||||
if len(values) == 0 {
|
||||
return nil
|
||||
}
|
||||
|
||||
budget := &ComputeBudget{
|
||||
MaxDailyHours: floatValue(values["max_daily_hours"]),
|
||||
MaxWeeklyCostUSD: floatValue(values["max_weekly_cost_usd"]),
|
||||
QuietStart: stringValue(values["quiet_start"]),
|
||||
QuietEnd: stringValue(values["quiet_end"]),
|
||||
PreferModels: listValue(values["prefer_models"]),
|
||||
AvoidModels: listValue(values["avoid_models"]),
|
||||
}
|
||||
|
||||
if computeBudgetIsZero(*budget) {
|
||||
return nil
|
||||
}
|
||||
return budget
|
||||
}
|
||||
|
||||
func computeBudgetIsZero(budget ComputeBudget) bool {
|
||||
return budget.MaxDailyHours == 0 &&
|
||||
budget.MaxWeeklyCostUSD == 0 &&
|
||||
core.Trim(budget.QuietStart) == "" &&
|
||||
core.Trim(budget.QuietEnd) == "" &&
|
||||
len(budget.PreferModels) == 0 &&
|
||||
len(budget.AvoidModels) == 0
|
||||
}
|
||||
|
||||
func parseFleetTask(values map[string]any) FleetTask {
|
||||
return FleetTask{
|
||||
ID: intValue(values["id"]),
|
||||
|
|
@ -1144,3 +1218,28 @@ func intValue(value any) int {
|
|||
}
|
||||
return 0
|
||||
}
|
||||
|
||||
func floatValue(value any) float64 {
|
||||
switch typed := value.(type) {
|
||||
case float64:
|
||||
return typed
|
||||
case float32:
|
||||
return float64(typed)
|
||||
case int:
|
||||
return float64(typed)
|
||||
case int64:
|
||||
return float64(typed)
|
||||
case string:
|
||||
trimmed := core.Trim(typed)
|
||||
if trimmed == "" {
|
||||
return 0
|
||||
}
|
||||
var parsed float64
|
||||
if result := core.JSONUnmarshalString(core.Concat("{\"n\":", trimmed, "}"), &struct {
|
||||
N *float64 `json:"n"`
|
||||
}{N: &parsed}); result.OK {
|
||||
return parsed
|
||||
}
|
||||
}
|
||||
return 0
|
||||
}
|
||||
|
|
|
|||
|
|
@ -71,6 +71,50 @@ func TestPlatform_HandleFleetRegister_Good(t *testing.T) {
|
|||
assert.Equal(t, []string{"go", "review"}, node.Capabilities)
|
||||
}
|
||||
|
||||
func TestPlatform_HandleFleetHeartbeat_Good_ComputeBudget(t *testing.T) {
|
||||
server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||||
require.Equal(t, "/v1/fleet/heartbeat", r.URL.Path)
|
||||
require.Equal(t, "Bearer secret-token", r.Header.Get("Authorization"))
|
||||
|
||||
bodyResult := core.ReadAll(r.Body)
|
||||
require.True(t, bodyResult.OK)
|
||||
|
||||
var payload map[string]any
|
||||
parseResult := core.JSONUnmarshalString(bodyResult.Value.(string), &payload)
|
||||
require.True(t, parseResult.OK)
|
||||
|
||||
budget, ok := payload["compute_budget"].(map[string]any)
|
||||
require.True(t, ok)
|
||||
assert.Equal(t, 2.5, budget["max_daily_hours"])
|
||||
assert.Equal(t, 12.5, budget["max_weekly_cost_usd"])
|
||||
assert.Equal(t, "22:00", budget["quiet_start"])
|
||||
assert.Equal(t, "06:00", budget["quiet_end"])
|
||||
assert.Equal(t, []any{"codex", "gpt-5.4"}, budget["prefer_models"])
|
||||
assert.Equal(t, []any{"gpt-4.1"}, budget["avoid_models"])
|
||||
|
||||
_, _ = w.Write([]byte(`{"data":{"id":1,"agent_id":"charon","platform":"linux","status":"online","compute_budget":{"max_daily_hours":2.5,"max_weekly_cost_usd":12.5,"quiet_start":"22:00","quiet_end":"06:00","prefer_models":["codex","gpt-5.4"],"avoid_models":["gpt-4.1"]}}}`))
|
||||
}))
|
||||
defer server.Close()
|
||||
|
||||
subsystem := testPrepWithPlatformServer(t, server, "secret-token")
|
||||
result := subsystem.handleFleetHeartbeat(context.Background(), core.NewOptions(
|
||||
core.Option{Key: "agent_id", Value: "charon"},
|
||||
core.Option{Key: "status", Value: "online"},
|
||||
core.Option{Key: "compute_budget", Value: `{"max_daily_hours":2.5,"max_weekly_cost_usd":12.5,"quiet_start":"22:00","quiet_end":"06:00","prefer_models":["codex","gpt-5.4"],"avoid_models":["gpt-4.1"]}`},
|
||||
))
|
||||
require.True(t, result.OK)
|
||||
|
||||
node, ok := result.Value.(FleetNode)
|
||||
require.True(t, ok)
|
||||
require.NotNil(t, node.ComputeBudget)
|
||||
assert.Equal(t, 2.5, node.ComputeBudget.MaxDailyHours)
|
||||
assert.Equal(t, 12.5, node.ComputeBudget.MaxWeeklyCostUSD)
|
||||
assert.Equal(t, "22:00", node.ComputeBudget.QuietStart)
|
||||
assert.Equal(t, "06:00", node.ComputeBudget.QuietEnd)
|
||||
assert.Equal(t, []string{"codex", "gpt-5.4"}, node.ComputeBudget.PreferModels)
|
||||
assert.Equal(t, []string{"gpt-4.1"}, node.ComputeBudget.AvoidModels)
|
||||
}
|
||||
|
||||
func TestPlatform_HandleFleetRegister_Bad(t *testing.T) {
|
||||
subsystem := testPrepWithPlatformServer(t, nil, "")
|
||||
|
||||
|
|
@ -291,7 +335,7 @@ func TestPlatform_HandleCreditsHistory_Good(t *testing.T) {
|
|||
|
||||
func TestPlatform_HandleFleetNodes_Good_NestedEnvelope(t *testing.T) {
|
||||
server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||||
_, _ = w.Write([]byte(`{"data":{"nodes":[{"id":1,"workspace_id":7,"agent_id":"charon","platform":"linux","models":["codex"],"status":"online"}],"total":1}}`))
|
||||
_, _ = w.Write([]byte(`{"data":{"nodes":[{"id":1,"workspace_id":7,"agent_id":"charon","platform":"linux","models":["codex"],"status":"online","compute_budget":{"max_daily_hours":3,"quiet_start":"22:00","prefer_models":["codex"]}}],"total":1}}`))
|
||||
}))
|
||||
defer server.Close()
|
||||
|
||||
|
|
@ -304,6 +348,10 @@ func TestPlatform_HandleFleetNodes_Good_NestedEnvelope(t *testing.T) {
|
|||
require.Len(t, output.Nodes, 1)
|
||||
assert.Equal(t, 1, output.Total)
|
||||
assert.Equal(t, 7, output.Nodes[0].WorkspaceID)
|
||||
require.NotNil(t, output.Nodes[0].ComputeBudget)
|
||||
assert.Equal(t, 3.0, output.Nodes[0].ComputeBudget.MaxDailyHours)
|
||||
assert.Equal(t, "22:00", output.Nodes[0].ComputeBudget.QuietStart)
|
||||
assert.Equal(t, []string{"codex"}, output.Nodes[0].ComputeBudget.PreferModels)
|
||||
}
|
||||
|
||||
func TestPlatform_HandleCreditsHistory_Good_NestedEnvelope(t *testing.T) {
|
||||
|
|
|
|||
|
|
@ -222,7 +222,7 @@ func (s *PrepSubsystem) fleetHeartbeatTool(ctx context.Context, _ *mcp.CallToolR
|
|||
options := platformOptions(
|
||||
core.Option{Key: "agent_id", Value: input.AgentID},
|
||||
core.Option{Key: "status", Value: input.Status},
|
||||
core.Option{Key: "compute_budget", Value: input.ComputeBudget},
|
||||
core.Option{Key: "compute_budget", Value: computeBudgetMapValue(input.ComputeBudget)},
|
||||
)
|
||||
result := s.handleFleetHeartbeat(ctx, options)
|
||||
if !result.OK {
|
||||
|
|
@ -235,6 +235,36 @@ func (s *PrepSubsystem) fleetHeartbeatTool(ctx context.Context, _ *mcp.CallToolR
|
|||
return nil, output, nil
|
||||
}
|
||||
|
||||
func computeBudgetMapValue(budget *ComputeBudget) map[string]any {
|
||||
if budget == nil || computeBudgetIsZero(*budget) {
|
||||
return nil
|
||||
}
|
||||
|
||||
values := map[string]any{}
|
||||
if budget.MaxDailyHours != 0 {
|
||||
values["max_daily_hours"] = budget.MaxDailyHours
|
||||
}
|
||||
if budget.MaxWeeklyCostUSD != 0 {
|
||||
values["max_weekly_cost_usd"] = budget.MaxWeeklyCostUSD
|
||||
}
|
||||
if trimmed := core.Trim(budget.QuietStart); trimmed != "" {
|
||||
values["quiet_start"] = trimmed
|
||||
}
|
||||
if trimmed := core.Trim(budget.QuietEnd); trimmed != "" {
|
||||
values["quiet_end"] = trimmed
|
||||
}
|
||||
if len(budget.PreferModels) > 0 {
|
||||
values["prefer_models"] = cleanStrings(budget.PreferModels)
|
||||
}
|
||||
if len(budget.AvoidModels) > 0 {
|
||||
values["avoid_models"] = cleanStrings(budget.AvoidModels)
|
||||
}
|
||||
if len(values) == 0 {
|
||||
return nil
|
||||
}
|
||||
return values
|
||||
}
|
||||
|
||||
func (s *PrepSubsystem) fleetDeregisterTool(ctx context.Context, _ *mcp.CallToolRequest, input FleetDeregisterInput) (*mcp.CallToolResult, map[string]any, error) {
|
||||
result := s.handleFleetDeregister(ctx, platformOptions(core.Option{Key: "agent_id", Value: input.AgentID}))
|
||||
if !result.OK {
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue