iop/cmd/iop-provider-smoke/main.go

package main

import (
	"context"
	"errors"
	"flag"
	"fmt"
	"os"
	"path/filepath"
	"strings"
	"sync"

	"go.uber.org/zap"

	"iop/packages/go/agentconfig"
	"iop/packages/go/agentprovider/catalog"
	runtime "iop/packages/go/agentruntime"
)

const smokeSessionID = "iop-provider-smoke"

type options struct {
	configPath string
	profileID  string
	operations string
	redact     bool
}

func main() {
	os.Exit(run(os.Args[1:]))
}

func run(arguments []string) int {
	var opts options
	flags := flag.NewFlagSet("iop-provider-smoke", flag.ContinueOnError)
	flags.SetOutput(os.Stderr)
	flags.StringVar(&opts.configPath, "config", "", "agent provider catalog YAML")
	flags.StringVar(&opts.profileID, "profile", "", "official profile ID")
	flags.StringVar(&opts.operations, "operations", "status,run,resume,cancel", "comma-separated lifecycle operations")
	flags.BoolVar(&opts.redact, "redact", false, "redact provider diagnostics")
	if err := flags.Parse(arguments); err != nil {
		return 2
	}
	if opts.configPath == "" || opts.profileID == "" {
		fmt.Fprintln(os.Stderr, "error=-config and -profile are required")
		return 2
	}
	if !opts.redact {
		fmt.Fprintln(os.Stderr, "error=-redact is required for provider smoke evidence")
		return 2
	}

	cfg, err := agentconfig.Load(opts.configPath)
	if err != nil {
		printError(err)
		return 1
	}
	discoverer, err := catalog.NewDiscoverer(cfg, nil)
	if err != nil {
		printError(err)
		return 1
	}
	readiness, err := discoverer.DiscoverProfile(context.Background(), opts.profileID)
	printPreflight(readiness)
	if err != nil {
		printError(err)
		return 1
	}

	provider, err := catalog.NewProfileProvider(cfg, opts.profileID, readiness, zap.NewNop())
	if err != nil {
		printError(err)
		return 1
	}
	defer func() {
		_ = provider.Stop(context.Background())
	}()

	workspace, err := os.Getwd()
	if err != nil {
		printError(err)
		return 1
	}
	for _, operation := range splitOperations(opts.operations) {
		if err := runOperation(context.Background(), provider, operation, workspace); err != nil {
			printError(fmt.Errorf("operation %s: %w", operation, err))
			return 1
		}
	}
	return 0
}

func printPreflight(readiness catalog.Readiness) {
	fmt.Printf(
		"preflight provider=%s model=%s profile=%s command=%s version=%s state=%s capabilities=%s redacted=true\n",
		readiness.ProviderID,
		readiness.ModelID,
		readiness.ProfileID,
		filepath.Base(readiness.Command),
		safeText(readiness.Version),
		readiness.State,
		strings.Join(readiness.Capabilities, ","),
	)
}

func runOperation(
	ctx context.Context,
	provider *catalog.ProfileProvider,
	operation, workspace string,
) error {
	switch operation {
	case "status":
		response, err := provider.HandleCommand(ctx, runtime.CommandRequest{
			RequestID: "smoke-status",
			Type:      runtime.CommandTypeUsageStatus,
		})
		if err != nil {
			return err
		}
		if response.UsageStatus == nil {
			return errors.New("status response is empty")
		}
		metadata := response.UsageStatus.Metadata
		fmt.Printf(
			"operation=status provider=%s model=%s profile=%s readiness=%s terminal=complete\n",
			metadata["provider_id"],
			metadata["model_id"],
			metadata["profile_id"],
			metadata["readiness"],
		)
		return nil
	case "run":
		return executeAndPrint(ctx, provider, runtime.ExecutionSpec{
			RunID:       "smoke-run",
			SessionID:   smokeSessionID,
			SessionMode: runtime.SessionModeCreateIfMissing,
			Workspace:   workspace,
			Input: map[string]any{
				"prompt": "Reply exactly IOP_PROVIDER_SMOKE_RUN_OK. Do not use tools or modify files.",
			},
		}, operation, runtime.EventTypeComplete, "IOP_PROVIDER_SMOKE_RUN_OK")
	case "resume":
		return executeAndPrint(ctx, provider, runtime.ExecutionSpec{
			RunID:       "smoke-resume",
			SessionID:   smokeSessionID,
			SessionMode: runtime.SessionModeRequireExisting,
			Workspace:   workspace,
			Input: map[string]any{
				"prompt": "Reply exactly IOP_PROVIDER_SMOKE_RESUME_OK. Do not use tools or modify files.",
			},
		}, operation, runtime.EventTypeComplete, "IOP_PROVIDER_SMOKE_RESUME_OK")
	case "cancel":
		cancelCtx, cancel := context.WithCancel(ctx)
		cancel()
		return executeAndPrint(cancelCtx, provider, runtime.ExecutionSpec{
			RunID:       "smoke-cancel",
			SessionID:   smokeSessionID,
			SessionMode: runtime.SessionModeCreateIfMissing,
			Workspace:   workspace,
			Input:       map[string]any{"prompt": "This invocation must be cancelled."},
		}, operation, runtime.EventTypeCancelled, "")
	default:
		return fmt.Errorf("unsupported operation %q", operation)
	}
}

func executeAndPrint(
	ctx context.Context,
	provider *catalog.ProfileProvider,
	spec runtime.ExecutionSpec,
	operation string,
	wantTerminal runtime.EventType,
	wantOutput string,
) error {
	sink := &smokeSink{}
	err := provider.Execute(ctx, spec, sink)
	if wantTerminal == runtime.EventTypeCancelled {
		if !errors.Is(err, runtime.ErrRunCancelled) {
			return fmt.Errorf("cancel error = %v, want %v", err, runtime.ErrRunCancelled)
		}
	} else if err != nil {
		return err
	}

	events := sink.Events()
	if len(events) == 0 {
		return errors.New("provider emitted no events")
	}
	terminal := events[len(events)-1]
	if terminal.Type != wantTerminal {
		return fmt.Errorf("terminal = %s, want %s", terminal.Type, wantTerminal)
	}
	for _, event := range events {
		if event.Metadata["provider_id"] == "" ||
			event.Metadata["model_id"] == "" ||
			event.Metadata["profile_id"] == "" {
			return fmt.Errorf("event %s is missing catalog identity", event.Type)
		}
	}
	if output := sink.Deltas(); wantOutput != "" && !strings.Contains(output, wantOutput) {
		return fmt.Errorf("output did not contain expected marker %q", wantOutput)
	}
	fmt.Printf(
		"operation=%s provider=%s model=%s profile=%s terminal=%s output=%s\n",
		operation,
		terminal.Metadata["provider_id"],
		terminal.Metadata["model_id"],
		terminal.Metadata["profile_id"],
		terminal.Type,
		safeText(sink.Deltas()),
	)
	return nil
}

type smokeSink struct {
	mu     sync.Mutex
	events []runtime.RuntimeEvent
}

func (s *smokeSink) Emit(_ context.Context, event runtime.RuntimeEvent) error {
	s.mu.Lock()
	defer s.mu.Unlock()
	s.events = append(s.events, event)
	return nil
}

func (s *smokeSink) Events() []runtime.RuntimeEvent {
	s.mu.Lock()
	defer s.mu.Unlock()
	return append([]runtime.RuntimeEvent(nil), s.events...)
}

func (s *smokeSink) Deltas() string {
	s.mu.Lock()
	defer s.mu.Unlock()
	var output strings.Builder
	for _, event := range s.events {
		if event.Type == runtime.EventTypeDelta {
			output.WriteString(event.Delta)
		}
	}
	return output.String()
}

func splitOperations(value string) []string {
	var operations []string
	for _, operation := range strings.Split(value, ",") {
		if trimmed := strings.TrimSpace(operation); trimmed != "" {
			operations = append(operations, trimmed)
		}
	}
	return operations
}

func safeText(value string) string {
	const maxOutput = 512
	value = catalog.Redact(strings.TrimSpace(value))
	value = strings.ReplaceAll(value, "\n", `\n`)
	value = strings.ReplaceAll(value, "\r", `\r`)
	if len(value) > maxOutput {
		return value[:maxOutput] + "…"
	}
	if value == "" {
		return "-"
	}
	return value
}

func printError(err error) {
	fmt.Fprintf(os.Stderr, "error=%s\n", safeText(err.Error()))
}