diff --git a/pkg/agentic/platform.go b/pkg/agentic/platform.go index 723568f..b40397c 100644 --- a/pkg/agentic/platform.go +++ b/pkg/agentic/platform.go @@ -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 +} diff --git a/pkg/agentic/platform_test.go b/pkg/agentic/platform_test.go index 4dc89d0..9d62219 100644 --- a/pkg/agentic/platform_test.go +++ b/pkg/agentic/platform_test.go @@ -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) { diff --git a/pkg/agentic/platform_tools.go b/pkg/agentic/platform_tools.go index 6b34c7c..9028592 100644 --- a/pkg/agentic/platform_tools.go +++ b/pkg/agentic/platform_tools.go @@ -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 {