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.
// ─────────────────────────────────────────────────────────────────────────────

// amStartArgs is the input to start a template agent. The template definition
// (.md) is maintained by HM and sent to AM here.
type amStartArgs struct {
	ManagerDeploymentID string
	UserID              string
	TemplateKey         string
	AgentDefinition     string // the agent .md (system prompt + frontmatter)
	Model               string
	Env                 map[string]string
	CallbackURL         string
}

// publicV1BaseURL is the model gateway base URL the started agent should call
// for models (OPENAI_BASE_URL).
func publicV1BaseURL() string {
	base := strings.TrimRight(strings.TrimSpace(common.GetEnvOrDefaultString("HEICODE_PUBLIC_BASE_URL", "https://code.xinghanlab.com")), "/")
	return base + "/v1"
}

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

func agentTemplateStartPath() string {
	// AM's agent-manager create endpoint (CODING_A2A doc §2: POST /agents).
	return common.GetEnvOrDefaultString("AGENT_RUNTIME_AGENT_START_PATH", "/agents")
}

func agentTemplateAgentPath(agentRuntimeID string) string {
	// Lifecycle lives under the same namespace as create (POST /agents).
	p := common.GetEnvOrDefaultString("AGENT_RUNTIME_AGENT_PATH", "/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", "/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
}

// amStartTemplateAgent asks AM to start a coding_a2a_agent. It maps HM's model
// (template .md + resource env) onto AM's create contract (CODING_A2A doc):
// POST /agents { name, template:"coding_a2a_agent", framework:"A2A",
// config:{user_id}, env:{...} } where the role/instruction/model are folded into
// env. AM's create response carries access_info.domain / namespace etc.
func amStartTemplateAgent(ctx context.Context, args amStartArgs) (amStartResult, error) {
	// Build AM's env: resource env + role/instruction + model gateway.
	env := map[string]string{}
	for k, v := range args.Env {
		env[k] = v
	}
	if args.TemplateKey != "" {
		env["AGENT_ROLE_NAME"] = args.TemplateKey
	}
	if strings.TrimSpace(args.AgentDefinition) != "" {
		env["AGENT_INSTRUCTION_TEXT"] = args.AgentDefinition
	}
	// MODEL_NAME must be a real model on the HM gateway. Template frontmatter
	// carries a Claude-style tier hint (opus/sonnet/haiku) which is NOT a gateway
	// model — fall back to the gateway default (gpt-5.4, env-overridable).
	modelName := common.GetEnvOrDefaultString("AGENT_RUNTIME_DEFAULT_MODEL", "gpt-5.4")
	switch strings.ToLower(strings.TrimSpace(args.Model)) {
	case "", "opus", "sonnet", "haiku", "claude":
		// keep the gateway default
	default:
		modelName = strings.TrimSpace(args.Model)
	}
	env["MODEL_NAME"] = modelName
	env["OPENAI_BASE_URL"] = publicV1BaseURL()
	// NOTE: OPENAI_API_KEY and AGENT_ACCESS_TOKEN are already in args.Env (the
	// handler mints + injects them before calling us) and are forwarded as-is by
	// the copy loop above. We don't add model/auth keys here.

	// env may carry plaintext secrets (db password, blob key…) — warn on http.
	base := agentRuntimeClientConfigForMode(agentRuntimeModeAgent).BaseURL
	if strings.HasPrefix(base, "http://") &&
		!strings.Contains(base, "localhost") && !strings.Contains(base, "127.0.0.1") {
		common.SysLog("WARNING: starting template agent with secret env over a non-HTTPS AM URL; use HTTPS or a private network")
	}

	name := strings.ToLower(strings.ReplaceAll(args.ManagerDeploymentID, "_", "-"))
	payload := map[string]any{
		"name":      name,
		"template":  "coding_a2a_agent",
		"framework": "A2A",
		"config": map[string]any{
			"user_id":               args.UserID,
			"manager_deployment_id": args.ManagerDeploymentID,
			"callback_url":          args.CallbackURL,
		},
		"env": env,
	}
	data, err := amTemplateDo(ctx, http.MethodPost, agentTemplateStartPath(), payload, agentTemplateStartTimeout())
	if err != nil {
		return amStartResult{}, err
	}
	// AM exposes the agent via access_info.{domain,external_ip}; the instance id
	// is name/namespace. AM's access_token (if any) is ignored — HM mints its own
	// per-agent access token (AGENT_ACCESS_TOKEN) and returns that to the client;
	// AM authorizes callers by comparing the X-Agent-Access-Token header to its
	// env AGENT_ACCESS_TOKEN locally.
	subdomain := stringFromMap(data, "subdomain", "address", "url")
	if subdomain == "" {
		if ai, ok := data["access_info"].(map[string]any); ok {
			subdomain = stringFromMap(ai, "domain", "external_ip", "url")
		}
	}
	return amStartResult{
		RuntimeID:   firstNonEmpty(stringFromMap(data, "runtime_id", "agent_id", "id", "name", "namespace"), name),
		Subdomain:   subdomain,
		AccessToken: stringFromMap(data, "access_token", "token"),
		Status:      firstNonEmpty(stringFromMap(data, "status", "runtime_status"), "running"),
	}, 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
}