main.go

// Code generated by ADL CLI v0.62.8. DO NOT EDIT.
// This file was automatically generated from an ADL (Agent Definition Language) specification.
// Manual changes to this file may be overwritten during regeneration.
//
// EXCEPTION: this file is intentionally listed in .adl-ignore and kept in sync
// with the generated template BY HAND. It carries exactly ONE deviation from
// `adl generate` output: the LLM client is the in-repo mock
// (internal/mock.NewMockLLMClient) instead of the generator's default
// OpenAI-compatible client, because mock-agent runs without real LLM
// credentials. Everything else must track the template - when re-syncing after
// an ADL CLI bump, regenerate into a scratch dir and re-apply only the mock
// client swap marked below. See the note next to `main.go` in .adl-ignore.

package main

import (
	"bytes"
	"context"
	"fmt"
	"log"
	"os"
	"os/signal"
	"path/filepath"
	"sort"
	"strings"
	"syscall"

	server "github.com/inference-gateway/adk/server"
	envconfig "github.com/sethvargo/go-envconfig"
	cobra "github.com/spf13/cobra"
	zap "go.uber.org/zap"
	yaml "gopkg.in/yaml.v3"

	config "github.com/inference-gateway/mock-agent/config"
	tools "github.com/inference-gateway/mock-agent/tools"

	logger "github.com/inference-gateway/mock-agent/internal/logger"
	mock "github.com/inference-gateway/mock-agent/internal/mock"
)

// Version, AgentName and AgentDescription are injected at build time
// via `-ldflags "-X 'main.Version=...'"` (see Dockerfile). They default
// to the values declared in the ADL.
var (
	Version          = "0.3.1"
	AgentName        = "mock-agent"
	AgentDescription = "A2A agent server for mocking and testing. Uses a mock LLM client - no API keys required!"
)

// skillsDir is the directory the runtime scans for skill manifests at
// startup. Override with A2A_SKILLS_DIR.
const skillsDir = ".agents/skills"

// defaultMaxChatCompletionIterations mirrors the ADK's own default and is used
// when the configured value is unset/0, so the agent runs multi-step tool-call
// workloads instead of stalling immediately.
const defaultMaxChatCompletionIterations = 50

// loadSkillsManifest walks the configured skills directory, reads each
// <skill>/SKILL.md, extracts the YAML frontmatter (name + description),
// and returns an `AVAILABLE SKILLS:` block to append to the system
// prompt. SKILL.md bodies are NOT inlined - the model must call the
// Read tool to load a skill's playbook on demand. Returns "" when the
// directory is missing or has no valid manifests.
func loadSkillsManifest(dir string) (string, error) {
	entries, err := os.ReadDir(dir)
	if err != nil {
		if os.IsNotExist(err) {
			return "", nil
		}
		return "", err
	}

	dirs := make([]string, 0, len(entries))
	for _, e := range entries {
		if !e.IsDir() {
			continue
		}
		dirs = append(dirs, e.Name())
	}
	sort.Strings(dirs)

	type skillFrontmatter struct {
		Name        string `yaml:"name"`
		Description string `yaml:"description"`
	}

	var manifest strings.Builder
	for _, id := range dirs {
		path := filepath.Join(dir, id, "SKILL.md")
		data, err := os.ReadFile(path)
		if err != nil {
			if os.IsNotExist(err) {
				continue
			}
			return "", err
		}
		fmBytes, ok := extractFrontmatter(data)
		if !ok {
			continue
		}
		var fm skillFrontmatter
		if err := yaml.Unmarshal(fmBytes, &fm); err != nil {
			continue
		}
		if fm.Name == "" || fm.Description == "" {
			continue
		}
		if manifest.Len() == 0 {
			manifest.WriteString("AVAILABLE SKILLS:\n")
			manifest.WriteString("Skills are reusable instructions for specific tasks. When a task matches a\n")
			manifest.WriteString("skill's description, read the SKILL.md file at the listed path using the Read\n")
			manifest.WriteString("tool, then follow its instructions.\n\n")
		}
		fmt.Fprintf(&manifest, "- %s: %s\n  Path: %s\n", fm.Name, fm.Description, path)
	}
	return manifest.String(), nil
}

// extractFrontmatter returns the bytes between the opening and closing
// `---` fences of a SKILL.md file (without including the fences
// themselves). The second return value is false when no frontmatter is
// found.
func extractFrontmatter(content []byte) ([]byte, bool) {
	bom := []byte{0xEF, 0xBB, 0xBF}
	buf := bytes.TrimPrefix(content, bom)
	buf = bytes.TrimLeft(buf, "\r\n\t ")
	if !bytes.HasPrefix(buf, []byte("---")) {
		return nil, false
	}
	rest := buf[3:]
	rest = bytes.TrimLeft(rest, "\r\n")
	idx := bytes.Index(rest, []byte("\n---"))
	if idx < 0 {
		return nil, false
	}
	return rest[:idx], true
}

// newRootCmd builds the top-level CLI for the agent binary. The
// generated binary is a real CLI: `<bin> --version`, `<bin> --help`,
// and `<bin> start` are all supported. Subcommands are added in
// dedicated constructors so they can be unit-tested in isolation.
func newRootCmd() *cobra.Command {
	root := &cobra.Command{
		Use:           AgentName,
		Short:         AgentDescription,
		Long:          AgentDescription + "\n\nThis is an A2A (Agent-to-Agent) protocol server. Use the `start` subcommand to run it.",
		Version:       Version,
		SilenceUsage:  true,
		SilenceErrors: true,
	}
	root.AddCommand(newStartCmd())
	return root
}

// newStartCmd returns the `start` subcommand which boots the A2A
// server and blocks until SIGINT/SIGTERM.
func newStartCmd() *cobra.Command {
	return &cobra.Command{
		Use:   "start",
		Short: "Start the A2A server",
		Long:  "Start the A2A server and block until SIGINT or SIGTERM is received.",
		RunE: func(cmd *cobra.Command, args []string) error {
			return runStart(cmd.Context())
		},
	}
}

// runStart contains the original agent bootstrap. It is exported as a
// dedicated function so the cobra command stays a thin shell - easier
// to test, easier to embed.
func runStart(ctx context.Context) error {
	var cfg config.Config
	if err := envconfig.Process(ctx, &cfg); err != nil {
		return fmt.Errorf("failed to load config: %w", err)
	}

	// AgentName and AgentVersion are build-time metadata (ldflags) that the ADK
	// deliberately does not map from the environment. Propagate them so the
	// OpenTelemetry resource carries service.name / service.version instead of
	// empty strings, and so any other consumer of cfg.A2A sees the real values.
	cfg.A2A.AgentName = AgentName
	cfg.A2A.AgentVersion = Version
	// The OpenTelemetry SDK settings are read as A2A_OTEL_* through the ADK's
	// A2A_-prefixed config (cfg.A2A.OTelConfig), so the single Process call above
	// already loaded them - no separate OTel pass is required.

	// Initialize logger
	l, err := logger.NewLogger(ctx, &cfg)
	if err != nil {
		return fmt.Errorf("failed to initialize logger: %w", err)
	}

	l.Info("starting "+AgentName+" agent", zap.String("version", Version), zap.Bool("debug", cfg.A2A.Debug))
	l.Debug("loaded configuration", zap.Any("config", cfg))

	resolvedSkillsDir := skillsDir
	if v := os.Getenv("A2A_SKILLS_DIR"); v != "" {
		resolvedSkillsDir = v
	}
	skillsPrompt, err := loadSkillsManifest(resolvedSkillsDir)
	if err != nil {
		l.Warn("failed to load skills manifest, continuing without them", zap.Error(err))
	} else if skillsPrompt != "" {
		l.Info("loaded skills manifest into system prompt", zap.String("dir", resolvedSkillsDir))
	}

	// Create toolbox with default tools (like input_required, create_artifact etc)
	toolBox := server.NewDefaultToolBox(&cfg.A2A.AgentConfig.ToolBoxConfig)

	// Register Read built-in
	readTool, err := tools.NewReadTool(ctx, l)
	if err != nil {
		return fmt.Errorf("failed to construct Read tool: %w", err)
	}
	toolBox.AddTool(readTool)
	l.Info("registered built-in: Read")

	// Register echo tool
	echoTool := tools.NewEchoTool()
	toolBox.AddTool(echoTool)
	l.Info("registered tool: echo (Echo back the input message (useful for basic connectivity tests))")

	// Register delay tool
	delayTool := tools.NewDelayTool()
	toolBox.AddTool(delayTool)
	l.Info("registered tool: delay (Simulate slow responses with configurable delays)")

	// Register error tool
	errorTool := tools.NewErrorTool()
	toolBox.AddTool(errorTool)
	l.Info("registered tool: error (Simulate error conditions for testing error handling)")

	// Register random_data tool
	randomDataTool := tools.NewRandomDataTool()
	toolBox.AddTool(randomDataTool)
	l.Info("registered tool: random_data (Generate random test data)")

	// Register validate tool
	validateTool := tools.NewValidateTool()
	toolBox.AddTool(validateTool)
	l.Info("registered tool: validate (Validate input against common patterns)")

	// Register simulate_tool_call tool
	simulateToolCallTool := tools.NewSimulateToolCallTool()
	toolBox.AddTool(simulateToolCallTool)
	l.Info("registered tool: simulate_tool_call (Simulate a single tool call for load, latency and failure testing. Emits an instrumented span (gen_ai.tool.name) with a configurable duration and optional error status. The mock LLM drives this once per entry of a multi-tool-call workload.)")

	// --- BEGIN mock deviation from the generated template -------------------
	// The template emits:
	//     llmClient, err := server.NewOpenAICompatibleLLMClient(&cfg.A2A.AgentConfig, l)
	//     if err != nil { return fmt.Errorf("failed to create LLM client: %w", err) }
	// mock-agent's whole purpose is to run without real LLM credentials, so we
	// substitute the in-repo mock client instead. This is the ONLY hand-applied
	// change relative to `adl generate`; keep it minimal so re-syncs stay trivial.
	llmClient := mock.NewMockLLMClient(l)
	l.Info("using mock LLM client (no external API calls)")
	// --- END mock deviation -------------------------------------------------

	systemPrompt := `You are a mock AI assistant designed for testing and development purposes.

You have access to several mock tools that demonstrate different testing scenarios:
- echo: Simply echo back the input message (useful for basic connectivity tests)
- delay: Simulate slow responses with configurable delays
- error: Simulate error conditions for testing error handling
- random_data: Generate random test data
- validate: Validate input against common patterns
- Read: Read a file from disk - the mock routes here when a request mentions
  "read <path>", exercising a real tool call (and its telemetry span) so
  distributed traces show a nested sub-tool span under a2a.request
- simulate_tool_call: Simulate one tool call with a configurable name, latency
  and optional error status. The mock drives it repeatedly to build multi
  tool-call workloads for load, latency and failure testing.

The mock routes on deterministic keywords, not model reasoning. Notably,
"read <path>" runs the Read tool against <path> (defaulting to README.md),
which is handy for end-to-end distributed-tracing demos. Saying
"simulate N tool calls" (or setting MOCK_TOOL_CALLS=read,search,read) drives a
sequence of instrumented simulate_tool_call spans - each with its own
gen_ai.tool.name, duration and optional injected failure - so a task can look
like a realistic multi-step agent in a trace.

When responding:
- Be clear and predictable in your responses
- Include relevant metadata about the request
- Support both streaming and non-streaming modes
- Handle edge cases gracefully

Your purpose is to provide consistent, reproducible responses for testing A2A protocol implementations.
`
	if skillsPrompt != "" {
		systemPrompt = systemPrompt + "\n\n" + skillsPrompt
	}

	maxIterations := cfg.A2A.AgentConfig.MaxChatCompletionIterations
	if maxIterations < 1 {
		maxIterations = defaultMaxChatCompletionIterations
	}

	agent, err := server.NewAgentBuilder(l).
		WithConfig(&cfg.A2A.AgentConfig).
		WithLLMClient(llmClient).
		WithToolBox(toolBox).
		WithMaxChatCompletion(maxIterations).
		WithSystemPrompt(systemPrompt).
		Build()
	if err != nil {
		return fmt.Errorf("failed to create agent: %w", err)
	}

	artifactService, err := server.NewArtifactService(&cfg.A2A.ArtifactsConfig, l)
	if err != nil {
		l.Warn("artifact service could not be created - check A2A_ARTIFACTS_ENABLE environment variable", zap.Error(err))
		l.Info("continuing without artifact service support")
		artifactService = nil
	}

	artifactsServer, err := server.
		NewArtifactsServerBuilder(&cfg.A2A.ArtifactsConfig, l).
		WithArtifactService(artifactService).
		Build()
	if err != nil {
		l.Warn("artifacts server could not be created", zap.Error(err))
		l.Info("continuing without artifacts server")
		artifactsServer = nil
	}

	a2aServer, err := server.NewA2AServerBuilder(cfg.A2A, l).
		WithAgent(agent).
		WithAgentCardFromFile(".well-known/agent-card.json", map[string]any{
			"name":        AgentName,
			"version":     Version,
			"description": AgentDescription,
			"url":         cfg.A2A.AgentURL,
		}).
		WithArtifactService(artifactService).
		WithDefaultBackgroundTaskHandler().
		WithDefaultStreamingTaskHandler().
		Build()
	if err != nil {
		return fmt.Errorf("failed to create A2A server: %w", err)
	}

	go func() {
		l.Info("starting A2A server", zap.String("port", cfg.A2A.ServerConfig.Port))
		if err := a2aServer.Start(ctx); err != nil {
			l.Fatal("server failed to start", zap.Error(err))
		}
	}()

	if artifactsServer != nil {
		go func() {
			l.Info("starting A2A artifacts server", zap.String("port", cfg.A2A.ArtifactsConfig.ServerConfig.Port))
			if err := artifactsServer.Start(ctx); err != nil {
				l.Fatal("artifacts server failed to start", zap.Error(err))
			}
		}()
	}

	l.Info("mock-agent agent running successfully",
		zap.String("port", cfg.A2A.ServerConfig.Port))

	quit := make(chan os.Signal, 1)
	signal.Notify(quit, syscall.SIGINT, syscall.SIGTERM)
	<-quit

	l.Info("shutdown signal received, gracefully stopping server...")
	a2aServer.Stop(ctx)
	if artifactsServer != nil {
		artifactsServer.Stop(ctx)
	}
	l.Info("mock-agent agent stopped")
	return nil
}

func main() {
	ctx := context.Background()
	if err := newRootCmd().ExecuteContext(ctx); err != nil {
		log.Fatal(err)
	}
}