/internal/harness/plan-delegate/main.go

// Plan & Delegate integration harness.
//
// This runs the REAL go-micro stack end to end — real services, real
// registry, real RPC, the real agent loop, real store, real delegate
// routing — and mocks ONLY the LLM with a deterministic provider. It
// proves the plumbing works without an API key, and it's reproducible.
//
// Swap MICRO_AI_PROVIDER/MICRO_AI_API_KEY (and remove --mock) to run the
// exact same flow against a live model.
//
// Run:
//
//	go run ./internal/harness/plan-delegate
package main

import (
	"context"
	"encoding/json"
	"flag"
	"fmt"
	"os"
	"strings"
	"sync"
	"time"

	"go-micro.dev/v6/agent"
	"go-micro.dev/v6/ai"
	"go-micro.dev/v6/broker"
	"go-micro.dev/v6/flow"
	"go-micro.dev/v6/internal/harness/harnessutil"
	"go-micro.dev/v6/registry"
	"go-micro.dev/v6/service"
	"go-micro.dev/v6/store"
)

// ---------------------------------------------------------------------------
// real services
// ---------------------------------------------------------------------------

type Task struct {
	ID    string `json:"id"`
	Title string `json:"title"`
}

type AddRequest struct {
	Title string `json:"title" description:"Title of the task to add"`
}
type AddResponse struct {
	Task *Task `json:"task"`
}
type ListRequest struct{}
type ListResponse struct {
	Tasks []*Task `json:"tasks"`
}

type TaskService struct {
	mu      sync.Mutex
	tasks   []*Task
	byTitle map[string]*Task
	nextID  int
}

// Add creates a new task with the given title. Replayed live-model tool calls
// are idempotent by launch task title so the conformance harness proves exactly
// one durable side effect per intended task even if a provider resends a call.
// @example {"title": "Design"}
func (s *TaskService) Add(ctx context.Context, req *AddRequest, rsp *AddResponse) error {
	s.mu.Lock()
	defer s.mu.Unlock()
	if s.byTitle == nil {
		s.byTitle = map[string]*Task{}
	}
	key := launchTaskKey(req.Title)
	if t := s.byTitle[key]; t != nil {
		rsp.Task = t
		fmt.Printf("    \033[32m[task]\033[0m reused %s %q\n", t.ID, t.Title)
		return nil
	}
	s.nextID++
	t := &Task{ID: fmt.Sprintf("task-%d", s.nextID), Title: canonicalLaunchTitle(req.Title)}
	s.tasks = append(s.tasks, t)
	s.byTitle[key] = t
	rsp.Task = t
	fmt.Printf("    \033[32m[task]\033[0m created %s %q\n", t.ID, t.Title)
	return nil
}

func launchTaskKey(title string) string {
	s := strings.ToLower(strings.TrimSpace(title))
	switch {
	case strings.Contains(s, "design"):
		return "design"
	case strings.Contains(s, "build"):
		return "build"
	case strings.Contains(s, "ship"):
		return "ship"
	default:
		return s
	}
}

func canonicalLaunchTitle(title string) string {
	switch launchTaskKey(title) {
	case "design":
		return "Design"
	case "build":
		return "Build"
	case "ship":
		return "Ship"
	default:
		return strings.TrimSpace(title)
	}
}

// List returns all tasks.
// @example {}
func (s *TaskService) List(ctx context.Context, req *ListRequest, rsp *ListResponse) error {
	s.mu.Lock()
	defer s.mu.Unlock()
	rsp.Tasks = append(rsp.Tasks, s.tasks...)
	return nil
}

func (s *TaskService) count() int {
	s.mu.Lock()
	defer s.mu.Unlock()
	return len(s.tasks)
}

const delegatedNotifyTask = "Use the notify Send tool exactly once to tell owner@acme.com: The launch plan is ready. Do not answer until the notify tool call has succeeded."

const commsPrompt = "You handle outbound notifications. When asked to notify someone, you must call the notify Send tool exactly once before replying. Never claim a notification was sent unless the notify tool returned success."

const delegatedNotifySettleTimeout = 10 * time.Second

type SendRequest struct {
	To      string `json:"to" description:"Recipient address"`
	Message string `json:"message" description:"Message body"`
}
type SendResponse struct {
	Sent bool `json:"sent"`
}
type NotifyService struct {
	mu         sync.Mutex
	sent       int
	attempts   int
	duplicates int
	bySend     map[string]bool
}

// Send delivers a notification message to a recipient. Duplicate delivery
// attempts for the same recipient/message are treated as successful replays
// without producing another side effect.
// @example {"to": "owner@acme.com", "message": "ready"}
func (s *NotifyService) Send(ctx context.Context, req *SendRequest, rsp *SendResponse) error {
	s.mu.Lock()
	if s.bySend == nil {
		s.bySend = map[string]bool{}
	}
	key := notifyDedupKey(req.To, req.Message)
	s.attempts++
	if !s.bySend[key] {
		s.bySend[key] = true
		s.sent++
		fmt.Printf("    \033[35m[notify]\033[0m 📨 to=%s message=%q\n", req.To, req.Message)
	} else {
		s.duplicates++
		fmt.Printf("    \033[35m[notify]\033[0m reused to=%s message=%q\n", req.To, req.Message)
	}
	s.mu.Unlock()
	rsp.Sent = true
	return nil
}

func (s *NotifyService) count() int {
	s.mu.Lock()
	defer s.mu.Unlock()
	return s.sent
}

func (s *NotifyService) duplicateAttempts() int {
	s.mu.Lock()
	defer s.mu.Unlock()
	return s.duplicates
}

func notifyDedupKey(to, message string) string {
	recipient := canonicalLaunchNotifyRecipient(normalizeNotifyText(to))
	body := normalizeNotifyText(message)
	if isLaunchReadinessNotify(body) {
		body = "launch-readiness"
	}
	return recipient + "\x00" + body
}

func canonicalLaunchNotifyRecipient(recipient string) string {
	switch recipient {
	case "owner", "launch owner", "plan owner", "owner acme com", "owner@acme com", "owner @ acme com":
		return "owner@acme.com"
	default:
		if strings.Contains(recipient, "owner") && strings.Contains(recipient, "acme") {
			return "owner@acme.com"
		}
		return recipient
	}
}

func normalizeNotifyText(message string) string {
	message = strings.ToLower(strings.TrimSpace(message))
	message = strings.Map(func(r rune) rune {
		switch {
		case r >= 'a' && r <= 'z', r >= '0' && r <= '9':
			return r
		case r == '@':
			return r
		default:
			return ' '
		}
	}, message)
	return strings.Join(strings.Fields(message), " ")
}

func isLaunchReadinessNotify(message string) bool {
	return strings.Contains(message, "launch") &&
		strings.Contains(message, "plan") &&
		(strings.Contains(message, "ready") ||
			strings.Contains(message, "readiness") ||
			strings.Contains(message, "prepared") ||
			strings.Contains(message, "complete"))
}

// ---------------------------------------------------------------------------
// mock LLM provider — the ONLY fake. It "reasons" by simple heuristics
// over the tools it's offered and the system prompt it's given, calling
// the real tool handler exactly the way a real provider would.
// ---------------------------------------------------------------------------

type mockModel struct {
	opts ai.Options

	// unknownDelegateOnce makes the mock emit one provider-style, unavailable
	// delegate tool name before using the registered delegate tool. This mirrors
	// live providers that occasionally hallucinate a provider-specific tool while
	// still keeping the regression deterministic and keyless.
	unknownDelegateOnce    bool
	emittedUnknownDelegate bool

	// duplicateNotify makes the comms mock replay the same notification call.
	// The notify service should collapse that replay to one durable side effect.
	duplicateNotify bool

	// duplicateDelegate makes the conductor mock replay the same delegate call.
	// The delegate idempotency path should collapse that replay before it can
	// ask the delegated comms agent to notify twice.
	duplicateDelegate bool
}

func newMock(opts ...ai.Option) ai.Model {
	m := &mockModel{}
	_ = m.Init(opts...)
	return m
}

func newMockUnknownDelegate(opts ...ai.Option) ai.Model {
	m := &mockModel{unknownDelegateOnce: true}
	_ = m.Init(opts...)
	return m
}

func newMockDuplicateNotify(opts ...ai.Option) ai.Model {
	m := &mockModel{duplicateNotify: true}
	_ = m.Init(opts...)
	return m
}

func newMockDuplicateDelegate(opts ...ai.Option) ai.Model {
	m := &mockModel{duplicateDelegate: true}
	_ = m.Init(opts...)
	return m
}

func (m *mockModel) Init(opts ...ai.Option) error {
	for _, o := range opts {
		o(&m.opts)
	}
	return nil
}
func (m *mockModel) Options() ai.Options { return m.opts }
func (m *mockModel) String() string      { return "mock" }
func (m *mockModel) Stream(ctx context.Context, req *ai.Request, _ ...ai.GenerateOption) (ai.Stream, error) {
	return nil, fmt.Errorf("stream not supported by mock")
}

// findTool returns the safe name of the first offered tool whose name
// contains sub, or "" if none.
func findTool(tools []ai.Tool, sub string) string {
	for _, t := range tools {
		if strings.Contains(t.Name, sub) {
			return t.Name
		}
	}
	return ""
}

func (m *mockModel) call(who, name string, input map[string]any) {
	args, _ := json.Marshal(input)
	fmt.Printf("  \033[33m[%s]\033[0m → %s(%s)\n", who, name, args)
	if m.opts.ToolHandler != nil {
		m.opts.ToolHandler(context.Background(), ai.ToolCall{Name: name, Input: input})
	}
}

func (m *mockModel) Generate(ctx context.Context, req *ai.Request, _ ...ai.GenerateOption) (*ai.Response, error) {
	// Classify by the tools actually offered, not by prompt text:
	// the conductor has the task "Add" tool, comms has "Send".
	hasAdd := findTool(req.Tools, "Add") != ""
	hasSend := findTool(req.Tools, "Send") != ""

	switch {
	// comms agent: owns notify, has Send but not Add.
	case hasSend && !hasAdd:
		send := findTool(req.Tools, "Send")
		input := map[string]any{
			"to":      "owner@acme.com",
			"message": "The launch plan is ready",
		}
		m.call("comms", send, input)
		if m.duplicateNotify {
			m.call("comms", send, input)
		}
		return &ai.Response{Answer: "Notified owner@acme.com."}, nil

	// conductor: has the task Add tool — plan, create tasks, delegate.
	case hasAdd:
		if plan := findTool(req.Tools, "plan"); plan != "" {
			m.call("conductor", plan, map[string]any{
				"steps": []any{
					map[string]any{"task": "create Design task", "status": "pending"},
					map[string]any{"task": "create Build task", "status": "pending"},
					map[string]any{"task": "create Ship task", "status": "pending"},
					map[string]any{"task": "notify owner via comms", "status": "pending"},
				},
			})
		}
		if add := findTool(req.Tools, "Add"); add != "" {
			for _, title := range []string{"Design", "Build", "Ship"} {
				m.call("conductor", add, map[string]any{"title": title})
			}
		}
		if del := findTool(req.Tools, "delegate"); del != "" {
			if m.unknownDelegateOnce && !m.emittedUnknownDelegate {
				m.emittedUnknownDelegate = true
				m.call("conductor", "atlascloud_delegate", map[string]any{
					"task": delegatedNotifyTask,
					"to":   "comms",
				})
			} else {
				input := map[string]any{
					"task": delegatedNotifyTask,
					"to":   "comms",
				}
				m.call("conductor", del, input)
				if m.duplicateDelegate {
					m.call("conductor", del, input)
				}
			}
		}
		return &ai.Response{Answer: "Created Design, Build and Ship, and had comms notify the owner."}, nil

	// ephemeral sub-agent or anything else.
	default:
		return &ai.Response{Reply: "subtask handled"}, nil
	}
}

func providerKey(provider string) string {
	if v := os.Getenv("MICRO_AI_API_KEY"); v != "" {
		return v
	}
	env := map[string]string{
		"anthropic":  "ANTHROPIC_API_KEY",
		"openai":     "OPENAI_API_KEY",
		"gemini":     "GEMINI_API_KEY",
		"groq":       "GROQ_API_KEY",
		"mistral":    "MISTRAL_API_KEY",
		"together":   "TOGETHER_API_KEY",
		"atlascloud": "ATLASCLOUD_API_KEY",
	}[provider]
	return os.Getenv(env)
}

func runPlanDelegate(provider string) error {
	apiKey := ""
	switch provider {
	case "mock":
		ai.Register("mock", newMock)
	case "mock-unknown-delegate":
		ai.Register("mock-unknown-delegate", newMockUnknownDelegate)
	case "mock-duplicate-notify":
		ai.Register("mock-duplicate-notify", newMockDuplicateNotify)
	case "mock-duplicate-delegate":
		ai.Register("mock-duplicate-delegate", newMockDuplicateDelegate)
	default:
		apiKey = providerKey(provider)
		if apiKey == "" {
			fmt.Printf("no API key for provider %q — set MICRO_AI_API_KEY or the provider's key env\n", provider)
			return nil
		}
	}

	fmt.Printf("\n\033[1mPlan & Delegate — live integration harness (provider: %s)\033[0m\n", provider)
	fmt.Print("Real services, registry, RPC, agent loop, store, delegation.\n\n")

	reg := registry.NewMemoryRegistry()
	cl := harnessutil.Client(provider, reg)
	mem := store.NewMemoryStore()
	liveAgentOpts := harnessutil.AgentOptions(provider)
	commsCheckpoint := flow.StoreCheckpoint(mem, "agent-comms")
	conductorCheckpoint := flow.StoreCheckpoint(mem, "agent-conductor")

	// Real services.
	taskSvc := new(TaskService)
	task := service.New(service.Name("task"), service.Address("127.0.0.1:0"), service.Registry(reg), service.Client(cl))
	if err := task.Handle(taskSvc); err != nil {
		return fmt.Errorf("task handle: %w", err)
	}
	go task.Run()

	notifySvc := new(NotifyService)
	notify := service.New(service.Name("notify"), service.Address("127.0.0.1:0"), service.Registry(reg), service.Client(cl))
	if err := notify.Handle(notifySvc); err != nil {
		return fmt.Errorf("notify handle: %w", err)
	}
	go notify.Run()

	// Real comms agent (owns notify), registered so delegate reaches it over RPC.
	commsOpts := []agent.Option{
		agent.Name("comms"),
		agent.Address("127.0.0.1:0"),
		agent.Services("notify"),
		agent.Prompt(commsPrompt),
		agent.Provider(provider), agent.APIKey(apiKey),
		agent.WithRegistry(reg), agent.WithClient(cl), agent.WithStore(mem),
		agent.WithCheckpoint(commsCheckpoint),
	}
	commsOpts = append(commsOpts, liveAgentOpts...)
	comms := agent.New(commsOpts...)
	go comms.Run()
	defer comms.Stop()

	// Real conductor agent (owns task), registered so the flow can reach it over RPC.
	conductorOpts := []agent.Option{
		agent.Name("conductor"),
		agent.Address("127.0.0.1:0"),
		agent.Services("task"),
		agent.Prompt("You coordinate launch work. Plan first, create exactly one Design task, one Build task, and one Ship task, then delegate exactly one readiness notification to the \"comms\" agent. Do not create duplicate tasks and do not send notifications yourself."),
		agent.Provider(provider), agent.APIKey(apiKey),
		agent.WithRegistry(reg), agent.WithClient(cl), agent.WithStore(mem),
		agent.WithCheckpoint(conductorCheckpoint),
	}
	conductorOpts = append(conductorOpts, liveAgentOpts...)
	conductor := agent.New(conductorOpts...)
	go conductor.Run()
	defer conductor.Stop()

	fmt.Println("waiting for services + agents to register...")
	waitForService := func(name string) error {
		deadline := time.Now().Add(5 * time.Second)
		for time.Now().Before(deadline) {
			if svcs, err := reg.GetService(name); err == nil && len(svcs) > 0 && len(svcs[0].Nodes) > 0 {
				return nil
			}
			time.Sleep(20 * time.Millisecond)
		}
		return fmt.Errorf("service %q never registered", name)
	}
	for _, name := range []string{"task", "notify", "comms", "conductor"} {
		if err := waitForService(name); err != nil {
			return err
		}
	}

	f := flow.New("zero-to-hero",
		flow.Steps(
			flow.Step{Name: "conductor", Run: planDelegateConductorStep(conductor, taskSvc, notifySvc)},
			flow.Step{Name: "require-notify", Run: requireDelegatedNotifyStep(taskSvc, notifySvc, func(ctx context.Context) error {
				_, err := comms.Ask(ctx, "Send exactly one owner readiness notification now with this exact task: "+delegatedNotifyTask+" Use the notify service and do not answer until the notification has been sent.")
				return err
			})},
		),
		flow.WithCheckpoint(flow.StoreCheckpoint(mem, "flow-zero-to-hero")),
		flow.Timeout(harnessutil.LiveTimeout(provider)),
	)
	if err := f.Register(reg, broker.DefaultBroker, cl); err != nil {
		return fmt.Errorf("flow register: %w", err)
	}

	fmt.Print("\n\033[1m> flow:\033[0m services + agents + workflow + plan/delegate, no API key.\n\n")
	ctx, cancel := context.WithCancel(context.Background())
	defer cancel()
	executeDone := make(chan error, 1)
	go func() {
		executeDone <- f.Execute(ctx, "launch readiness")
	}()

	if err := waitForPlanDelegateExecution(executeDone, taskSvc, notifySvc); err != nil {
		return err
	}

	// Prove plan was persisted to the real store.
	if recs, _ := store.Scope(mem, "agent", "conductor").Read("plan"); len(recs) > 0 {
		fmt.Printf("\n\033[1mstored plan (agent/conductor/plan):\033[0m %s\n", string(recs[0].Value))
	} else {
		return fmt.Errorf("plan was not persisted")
	}
	if taskSvc.count() != 3 || notifySvc.count() != 1 {
		return fmt.Errorf("unexpected side effects: tasks=%d notify=%d", taskSvc.count(), notifySvc.count())
	}

	fmt.Println("\n\033[32m✓ 0→hero flow complete (services → agents → workflow)\033[0m")
	return nil
}

func planDelegateConductorStep(conductor agent.Agent, taskSvc *TaskService, notifySvc *NotifyService) flow.StepFunc {
	return func(ctx context.Context, in flow.State) (flow.State, error) {
		prompt := "Create three launch tasks (Design, Build, Ship), then make sure owner@acme.com is notified: " + in.String()
		rsp, err := conductor.Ask(ctx, prompt)
		if err != nil {
			if isUnfinishedPlanError(err) && taskSvc != nil && notifySvc != nil && taskSvc.count() == 3 && notifySvc.count() == 0 {
				fmt.Printf("\n\033[33mwarning:\033[0m conductor stopped with unfinished delegation after creating tasks; continuing to require-notify recovery: %v\n", err)
				return in, nil
			}
			return in, err
		}
		if rsp != nil && rsp.Reply != "" {
			fmt.Println("\n\033[1m< conductor reply:\033[0m", rsp.Reply)
		}
		return in, nil
	}
}

func isUnfinishedPlanError(err error) bool {
	if err == nil {
		return false
	}
	return strings.Contains(strings.ToLower(err.Error()), "unfinished plan steps")
}

func requireDelegatedNotifyStep(taskSvc *TaskService, notifySvc *NotifyService, recoverMissingNotify func(context.Context) error) flow.StepFunc {
	return func(ctx context.Context, in flow.State) (flow.State, error) {
		tasks := taskSvc.count()
		notify := notifySvc.count()
		if notify == 1 {
			return in, nil
		}
		if recoverMissingNotify == nil || tasks != 3 || notify != 0 {
			return in, fmt.Errorf("delegation completed without required notify side effect: notify=%d, want 1", notify)
		}
		settled, err := waitForNotifySideEffect(notifySvc, delegatedNotifySettleTimeout)
		if err != nil {
			return in, err
		}
		if !settled {
			fmt.Print("\n\033[33mwarning:\033[0m conductor step completed before delegated notify; retrying the missing comms handoff once before the flow can complete.\n")
			if err := recoverMissingNotify(ctx); err != nil {
				return in, fmt.Errorf("delegation completed without required notify side effect and recovery failed: notify=%d, want 1: %w", notify, err)
			}
			settled, err = waitForNotifySideEffect(notifySvc, delegatedNotifySettleTimeout)
			if err != nil {
				return in, err
			}
			if !settled {
				return in, fmt.Errorf("delegation recovery completed without required notify side effect: notify=%d, want 1", notifySvc.count())
			}
		}
		if notify = notifySvc.count(); notify != 1 {
			return in, fmt.Errorf("delegation recovery completed without required notify side effect: notify=%d, want 1", notify)
		}
		return in, nil
	}
}

func waitForPlanDelegateExecution(done <-chan error, taskSvc *TaskService, notifySvc *NotifyService) error {
	ticker := time.NewTicker(50 * time.Millisecond)
	defer ticker.Stop()
	for {
		select {
		case err := <-done:
			tasks := taskSvc.count()
			notify := notifySvc.count()
			if err != nil {
				if isClientTimeout(err) {
					if tasks == 3 && notify == 1 {
						fmt.Printf("\n\033[33mwarning:\033[0m flow execute returned after completed side effects: %v\n", err)
						return nil
					}
					return classifiedPlanDelegateTimeout(tasks, notify, err)
				}
				return fmt.Errorf("flow execute after side effects tasks=%d notify=%d: %w", tasks, notify, err)
			}
			if notify != 1 {
				return fmt.Errorf("delegation completed without required notify side effect: notify=%d, want 1", notify)
			}
			return nil
		case <-ticker.C:
			if notifySvc.count() == 1 {
				continue
			}
		}
	}
}

func waitForNotifySideEffect(notifySvc *NotifyService, timeout time.Duration) (bool, error) {
	deadline := time.Now().Add(timeout)
	for {
		if notifySvc.count() == 1 {
			return true, nil
		}
		if !time.Now().Before(deadline) {
			return false, nil
		}
		time.Sleep(50 * time.Millisecond)
	}
}

func classifiedPlanDelegateTimeout(tasks, notify int, err error) error {
	return fmt.Errorf("provider latency/outage during plan-delegate before required side effects completed (tasks=%d/3 notify=%d/1); retry live provider or inspect provider logs if this recurs: %w", tasks, notify, err)
}

func isClientTimeout(err error) bool {
	msg := strings.ToLower(err.Error())
	return strings.Contains(msg, "request timeout") || strings.Contains(msg, "code=408") || strings.Contains(msg, "code\":408")
}

func main() {
	provider := flag.String("provider", "mock", "LLM provider: mock (default), mock-unknown-delegate, mock-duplicate-notify, mock-duplicate-delegate, anthropic, openai, gemini, groq, mistral, together, atlascloud")
	flag.Parse()

	if err := runPlanDelegate(*provider); err != nil {
		fmt.Println("\033[31merror:\033[0m", err)
		os.Exit(1)
	}
}