heicode-mananger/heicode/controller/agent_template_runtime.go

package controller

import (
	"bytes"
	"context"
	"errors"
	"fmt"
	"io"
	"net/http"
	"net/url"
	"strings"
	"time"

	"github.com/heicode/manager/common"
)

// ─────────────────────────────────────────────────────────────────────────────
// AM template-agent adapter (the ONLY place that knows AM's template HTTP API).
//
// Template-agent model: HM asks AM to start a template agent with the selected
// resources injected as env; AM returns a unique subdomain + access token; the
// desktop client then talks to that subdomain directly over SSE.
//
// When AM finalises its real API, change ONLY:
//   - the default paths in the *Path() helpers (or set the env overrides), and
//   - the request payload / response field names in the three am* functions.
// No other HM code needs to change.
// ─────────────────────────────────────────────────────────────────────────────

// agentTemplate is one deployable template offered by AM.
type agentTemplate struct {
	TemplateID            string   `json:"template_id"`
	Name                  string   `json:"name"`
	Description           string   `json:"description"`
	RequiredResourceTypes []string `json:"required_resource_types"`
	EnvSchema             []string `json:"env_schema"`
}

// amStartResult is what AM returns after starting a template agent.
type amStartResult struct {
	RuntimeID   string
	Subdomain   string
	AccessToken string
	Status      string
}

func agentTemplatesPath() string {
	return common.GetEnvOrDefaultString("AGENT_RUNTIME_TEMPLATES_PATH", "/api/agent/templates")
}

func agentTemplateStartPath(templateID string) string {
	p := common.GetEnvOrDefaultString("AGENT_RUNTIME_TEMPLATE_START_PATH", "/api/agent/templates/{template_id}/start")
	return strings.ReplaceAll(p, "{template_id}", url.PathEscape(templateID))
}

func agentTemplateAgentPath(agentRuntimeID string) string {
	p := common.GetEnvOrDefaultString("AGENT_RUNTIME_AGENT_PATH", "/api/agent/agents/{agent_id}")
	return strings.ReplaceAll(p, "{agent_id}", url.PathEscape(agentRuntimeID))
}

func agentTemplateAgentStopPath(agentRuntimeID string) string {
	p := common.GetEnvOrDefaultString("AGENT_RUNTIME_AGENT_STOP_PATH", "/api/agent/agents/{agent_id}/stop")
	return strings.ReplaceAll(p, "{agent_id}", url.PathEscape(agentRuntimeID))
}

// agentTemplateStartTimeout is the timeout for starting a template agent —
// longer than the default runtime timeout because AM provisions the agent
// (container + subdomain) synchronously before responding.
func agentTemplateStartTimeout() time.Duration {
	sec := common.GetEnvOrDefault("AGENT_RUNTIME_START_TIMEOUT_SECONDS", 60)
	if sec <= 0 {
		sec = 60
	}
	return time.Duration(sec) * time.Second
}

// amTemplateDo performs an AM HTTP call and returns the decoded `data` object.
// Reuses the existing runtime config (base URL / service token). A non-zero
// timeout overrides the default runtime timeout (used by the start call).
func amTemplateDo(ctx context.Context, method, path string, body any, timeout time.Duration) (map[string]any, error) {
	cfg := agentRuntimeClientConfigForMode(agentRuntimeModeAgent)
	endpoint, err := agentRuntimeURL(cfg.BaseURL, path)
	if err != nil {
		return nil, err
	}
	if timeout <= 0 {
		timeout = cfg.Timeout
	}
	var reader io.Reader
	if body != nil {
		raw, err := common.Marshal(body)
		if err != nil {
			return nil, err
		}
		reader = bytes.NewReader(raw)
	}
	req, err := http.NewRequestWithContext(ctx, method, endpoint, reader)
	if err != nil {
		return nil, err
	}
	req.Header.Set("Content-Type", "application/json")
	if cfg.Token != "" {
		req.Header.Set("Authorization", "Bearer "+cfg.Token)
	}
	client := &http.Client{Timeout: timeout}
	resp, err := client.Do(req)
	if err != nil {
		return nil, err
	}
	defer resp.Body.Close()
	raw, _ := io.ReadAll(io.LimitReader(resp.Body, 1<<20))
	if resp.StatusCode < http.StatusOK || resp.StatusCode >= http.StatusMultipleChoices {
		return nil, fmt.Errorf("AM %s %s returned HTTP %d: %s", method, path, resp.StatusCode, strings.TrimSpace(string(raw)))
	}
	var envelope map[string]any
	if len(raw) > 0 {
		if err := common.Unmarshal(raw, &envelope); err != nil {
			return nil, err
		}
	}
	if msg := agentRuntimeEnvelopeError(envelope); msg != "" {
		return nil, errors.New(msg)
	}
	return extractAgentRuntimeData(envelope), nil
}

// amListTemplates lists AM's deployable templates.
// Proposed contract: data.templates = [...]. Swap this mapping if AM differs.
func amListTemplates(ctx context.Context) ([]agentTemplate, error) {
	data, err := amTemplateDo(ctx, http.MethodGet, agentTemplatesPath(), nil, 0)
	if err != nil {
		return nil, err
	}
	raw, err := common.Marshal(data["templates"])
	if err != nil {
		return nil, err
	}
	var templates []agentTemplate
	if err := common.Unmarshal(raw, &templates); err != nil {
		return nil, err
	}
	return templates, nil
}

// amStartTemplateAgent asks AM to start a template agent with the given env.
// Proposed request: {manager_deployment_id, env, callback_url}.
// Proposed response data: {runtime_id, subdomain, access_token, status}.
func amStartTemplateAgent(ctx context.Context, templateID, managerDeploymentID string, env map[string]string, callbackURL string) (amStartResult, error) {
	payload := map[string]any{
		"manager_deployment_id": managerDeploymentID,
		"env":                   env,
		"callback_url":          callbackURL,
	}
	data, err := amTemplateDo(ctx, http.MethodPost, agentTemplateStartPath(templateID), payload, agentTemplateStartTimeout())
	if err != nil {
		return amStartResult{}, err
	}
	return amStartResult{
		RuntimeID:   stringFromMap(data, "runtime_id", "agent_id", "id", "deployment_id"),
		Subdomain:   stringFromMap(data, "subdomain", "address", "url"),
		AccessToken: stringFromMap(data, "access_token", "token"),
		Status:      stringFromMap(data, "status", "runtime_status"),
	}, nil
}

// amGetAgentStatus pulls the live status of a running template agent from AM,
// so HM can tell whether the agent is still alive / has crashed.
// Proposed contract: GET /api/agent/agents/{agent_id} -> data.status.
func amGetAgentStatus(ctx context.Context, runtimeID string) (string, error) {
	data, err := amTemplateDo(ctx, http.MethodGet, agentTemplateAgentPath(runtimeID), nil, 0)
	if err != nil {
		return "", err
	}
	return stringFromMap(data, "status", "runtime_status", "state"), nil
}

func amStopTemplateAgent(ctx context.Context, runtimeID string) error {
	_, err := amTemplateDo(ctx, http.MethodPost, agentTemplateAgentStopPath(runtimeID), nil, 0)
	return err
}

func amDeleteTemplateAgent(ctx context.Context, runtimeID string) error {
	_, err := amTemplateDo(ctx, http.MethodDelete, agentTemplateAgentPath(runtimeID), nil, 0)
	return err
}