/internal/harness/a2a-stream-fallback/main.go

// A2A stream fallback harness.
//
// It exercises the gateway boundary that fronts an agent over A2A. The agent is
// configured with tools and memory, but its model streaming path deliberately
// reports ai.ErrStreamingUnsupported; the A2A gateway must fall back to the
// normal Ask path and still complete the same tool-calling run.
package main

import (
	"bufio"
	"bytes"
	"context"
	"encoding/json"
	"errors"
	"flag"
	"fmt"
	"io"
	"net/http"
	"net/http/httptest"
	"os"
	"strings"
	"time"

	"go-micro.dev/v6/agent"
	"go-micro.dev/v6/ai"
	"go-micro.dev/v6/gateway/a2a"
	"go-micro.dev/v6/registry"
	"go-micro.dev/v6/store"
)

type mockModel struct{ opts ai.Options }

func newMock(opts ...ai.Option) ai.Model {
	m := &mockModel{}
	_ = 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(context.Context, *ai.Request, ...ai.GenerateOption) (ai.Stream, error) {
	return nil, ai.ErrStreamingUnsupported
}
func (m *mockModel) Generate(ctx context.Context, req *ai.Request, _ ...ai.GenerateOption) (*ai.Response, error) {
	if req.Prompt == "" {
		return nil, errors.New("missing prompt")
	}
	if len(req.Messages) == 0 || req.Messages[len(req.Messages)-1].Role != "user" {
		return nil, fmt.Errorf("missing user history: %+v", req.Messages)
	}
	if len(req.Tools) == 0 || m.opts.ToolHandler == nil {
		return nil, errors.New("missing tools or tool handler")
	}
	res := m.opts.ToolHandler(ctx, ai.ToolCall{ID: "a2a-fallback-call", Name: "fallback_echo", Input: map[string]any{"value": "a2a-fallback"}})
	if res.Content == "" {
		return nil, errors.New("empty tool result")
	}
	return &ai.Response{Reply: "fallback completed", Answer: res.Content, ToolCalls: []ai.ToolCall{{ID: "a2a-fallback-call", Name: "fallback_echo", Input: map[string]any{"value": "a2a-fallback"}, Result: res.Content}}}, 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 main() {
	provider := flag.String("provider", "mock", "LLM provider: mock (default), anthropic, openai, ...")
	flag.Parse()

	apiKey := ""
	if *provider == "mock" {
		ai.Register("mock", newMock)
	} else {
		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
		}
	}

	fmt.Printf("\n\033[1mA2A streaming fallback conformance (provider: %s)\033[0m\n", *provider)
	reg := registry.NewMemoryRegistry()
	st := store.NewMemoryStore()
	var sawTool, sawRunInfo bool
	ag := agent.New(
		agent.Name("a2a-fallback"),
		agent.Provider(*provider),
		agent.APIKey(apiKey),
		agent.Prompt("Use fallback_echo exactly once with value a2a-fallback, then answer with the tool result."),
		agent.WithRegistry(reg),
		agent.WithStore(st),
		agent.WithMemory(agent.NewInMemory(8)),
		agent.ModelCallTimeout(45*time.Second),
		agent.WithTool("fallback_echo", "Echo the A2A fallback marker.", map[string]any{
			"value": map[string]any{"type": "string", "description": "value to echo"},
		}, func(ctx context.Context, input map[string]any) (string, error) {
			sawTool = true
			info, ok := ai.RunInfoFrom(ctx)
			if !ok || info.RunID == "" || info.Agent != "a2a-fallback" {
				return "", fmt.Errorf("unexpected run info: %+v", info)
			}
			sawRunInfo = true
			if input["value"] != "a2a-fallback" {
				return "", fmt.Errorf("unexpected value %v", input["value"])
			}
			return `{"marker":"a2a-fallback-ok"}`, nil
		}),
	)

	card := a2a.Card("a2a-fallback", "http://example.invalid/a2a-fallback", "", nil)
	handler := a2a.NewAgentStreamHandler(card, func(ctx context.Context, text string) (string, error) {
		resp, err := ag.Ask(ctx, text)
		if err != nil {
			return "", err
		}
		return resp.Reply, nil
	}, ag.Stream)

	body := []byte(`{"jsonrpc":"2.0","id":1,"method":"message/stream","params":{"message":{"role":"user","parts":[{"kind":"text","text":"Run the A2A fallback conformance check."}],"kind":"message"}}}`)
	req := httptest.NewRequest(http.MethodPost, "/", bytes.NewReader(body))
	req.Header.Set("Content-Type", "application/json")
	rr := httptest.NewRecorder()
	handler.ServeHTTP(rr, req)
	res := rr.Result()
	defer res.Body.Close()
	if res.StatusCode != http.StatusOK {
		b, _ := io.ReadAll(res.Body)
		fmt.Fprintf(os.Stderr, "unexpected status %d: %s\n", res.StatusCode, b)
		os.Exit(1)
	}
	if ct := res.Header.Get("Content-Type"); !strings.HasPrefix(ct, "text/event-stream") {
		fmt.Fprintf(os.Stderr, "content-type = %q, want text/event-stream\n", ct)
		os.Exit(1)
	}
	payload, err := readSSEData(res.Body)
	if err != nil {
		fmt.Fprintln(os.Stderr, err)
		os.Exit(1)
	}
	if !strings.Contains(payload, "a2a-fallback-ok") {
		fmt.Fprintf(os.Stderr, "stream payload missing marker: %s\n", payload)
		os.Exit(1)
	}
	if !sawTool || !sawRunInfo {
		fmt.Fprintf(os.Stderr, "tool=%v runInfo=%v\n", sawTool, sawRunInfo)
		os.Exit(1)
	}
	fmt.Println("\n\033[32m✓ A2A message/stream fell back to Ask and preserved tool/run metadata\033[0m")
}

func readSSEData(r io.Reader) (string, error) {
	scanner := bufio.NewScanner(r)
	var payload strings.Builder
	for scanner.Scan() {
		line := scanner.Text()
		if data, ok := strings.CutPrefix(line, "data: "); ok {
			payload.WriteString(data)
			payload.WriteByte('\n')
		}
	}
	if err := scanner.Err(); err != nil {
		return "", err
	}
	if payload.Len() == 0 {
		return "", errors.New("no SSE data received")
	}
	if !json.Valid([]byte(strings.TrimSpace(payload.String()))) {
		return "", fmt.Errorf("SSE data is not JSON: %s", payload.String())
	}
	return payload.String(), nil
}