diff --git a/.changeset/patch-add-model-aliases.md b/.changeset/patch-add-model-aliases.md new file mode 100644 index 00000000000..bc8fcf6f2b1 --- /dev/null +++ b/.changeset/patch-add-model-aliases.md @@ -0,0 +1,5 @@ +--- +"gh-aw": patch +--- + +Added model alias configuration support with built-in aliases and workflow import/frontmatter merging. diff --git a/pkg/parser/import_field_extractor.go b/pkg/parser/import_field_extractor.go index 664bd197a97..a3dd878e37f 100644 --- a/pkg/parser/import_field_extractor.go +++ b/pkg/parser/import_field_extractor.go @@ -50,7 +50,8 @@ type importAccumulator struct { skipBotsSet map[string]bool caches []string features []map[string]any - runInstallScripts bool // true if any imported workflow sets run-install-scripts: true (global or node-level) + models []map[string][]string // model alias maps from each imported file (appended in import order) + runInstallScripts bool // true if any imported workflow sets run-install-scripts: true (global or node-level) agentFile string agentImportSpec string repositoryImports []string @@ -436,6 +437,30 @@ func (acc *importAccumulator) extractAllImportFields(content []byte, item import } } + // Extract model aliases from imported file (parse as map[string][]string structure) + modelsContent, err := extractFieldJSONFromMap(fm, "models", "{}") + if err == nil && modelsContent != "" && modelsContent != "{}" { + var rawModels map[string]any + if jsonErr := json.Unmarshal([]byte(modelsContent), &rawModels); jsonErr == nil { + modelsMap := make(map[string][]string, len(rawModels)) + for k, v := range rawModels { + if patterns, ok := v.([]any); ok { + strs := make([]string, 0, len(patterns)) + for _, p := range patterns { + if s, ok := p.(string); ok { + strs = append(strs, s) + } + } + modelsMap[k] = strs + } + } + if len(modelsMap) > 0 { + acc.models = append(acc.models, modelsMap) + log.Printf("Extracted model aliases from import: %d entries", len(modelsMap)) + } + } + } + // Extract run-install-scripts flag from imported file. // If global run-install-scripts: true is set OR if runtimes.node.run-install-scripts: true is set, // propagate to the accumulator (OR semantics: any import enabling it enables it overall). @@ -510,6 +535,7 @@ func (acc *importAccumulator) toImportsResult(topologicalOrder []string) *Import MergedEnv: acc.envBuilder.String(), MergedEnvSources: acc.envSources, MergedFeatures: acc.features, + MergedModels: acc.models, MergedObservability: acc.observabilityBuilder.String(), ImportedFiles: topologicalOrder, AgentFile: acc.agentFile, diff --git a/pkg/parser/import_processor.go b/pkg/parser/import_processor.go index acbaf18ce7c..a7575e7281b 100644 --- a/pkg/parser/import_processor.go +++ b/pkg/parser/import_processor.go @@ -14,42 +14,43 @@ var importLog = logger.New("parser:import_processor") // ImportsResult holds the result of processing imports from frontmatter type ImportsResult struct { - MergedTools string // Merged tools configuration from all imports - MergedMCPServers string // Merged mcp-servers configuration from all imports - MergedEngines []string // Merged engine configurations from all imports - MergedSafeOutputs []string // Merged safe-outputs configurations from all imports - MergedMCPScripts []string // Merged mcp-scripts configurations from all imports - MergedMarkdown string // Only contains imports WITH inputs (for compile-time substitution) - ImportPaths []string // List of import file paths for runtime-import macro generation (replaces MergedMarkdown) - MergedSteps string // Merged steps configuration from all imports (excluding copilot-setup-steps) - CopilotSetupSteps string // Steps from copilot-setup-steps.yml (inserted at start) - MergedPreSteps string // Merged pre-steps configuration from all imports (prepended in order) - MergedPreAgentSteps string // Merged pre-agent-steps configuration from all imports (prepended in order) - MergedRuntimes string // Merged runtimes configuration from all imports - MergedRunInstallScripts bool // true if any imported workflow sets run-install-scripts: true (global or node-level) - MergedServices string // Merged services configuration from all imports - MergedNetwork string // Merged network configuration from all imports - MergedPermissions string // Merged permissions configuration from all imports - MergedSecretMasking string // Merged secret-masking steps from all imports - MergedBots []string // Merged bots list from all imports (union of bot names) - MergedSkipRoles []string // Merged skip-roles list from all imports (union of role names) - MergedSkipBots []string // Merged skip-bots list from all imports (union of usernames) - MergedActivationGitHubToken string // GitHub token from on.github-token in first imported workflow that defines it - MergedActivationGitHubApp string // JSON-encoded on.github-app from first imported workflow that defines it - MergedTopLevelGitHubApp string // JSON-encoded top-level github-app from first imported workflow that defines it - MergedCheckout string // JSON-encoded checkout configurations from imported workflows (one JSON value per line) - MergedPostSteps string // Merged post-steps configuration from all imports (appended in order) - MergedLabels []string // Merged labels from all imports (union of label names) - MergedCaches []string // Merged cache configurations from all imports (appended in order) - MergedJobs string // Merged jobs from imported YAML workflows (JSON format) - MergedEnv string // Merged env configuration from all imports (JSON format) - MergedEnvSources map[string]string // env var name → source import path (for conflict detection and lock file header listing) - MergedFeatures []map[string]any // Merged features configuration from all imports (parsed YAML structures) - MergedObservability string // Observability config (JSON) from first import that defines it (first-wins) - ImportedFiles []string // List of imported file paths (for manifest) - AgentFile string // Path to custom agent file (if imported) - AgentImportSpec string // Original import specification for agent file (e.g., "owner/repo/path@ref") - RepositoryImports []string // List of repository imports (format: "owner/repo@ref") for .github folder merging + MergedTools string // Merged tools configuration from all imports + MergedMCPServers string // Merged mcp-servers configuration from all imports + MergedEngines []string // Merged engine configurations from all imports + MergedSafeOutputs []string // Merged safe-outputs configurations from all imports + MergedMCPScripts []string // Merged mcp-scripts configurations from all imports + MergedMarkdown string // Only contains imports WITH inputs (for compile-time substitution) + ImportPaths []string // List of import file paths for runtime-import macro generation (replaces MergedMarkdown) + MergedSteps string // Merged steps configuration from all imports (excluding copilot-setup-steps) + CopilotSetupSteps string // Steps from copilot-setup-steps.yml (inserted at start) + MergedPreSteps string // Merged pre-steps configuration from all imports (prepended in order) + MergedPreAgentSteps string // Merged pre-agent-steps configuration from all imports (prepended in order) + MergedRuntimes string // Merged runtimes configuration from all imports + MergedRunInstallScripts bool // true if any imported workflow sets run-install-scripts: true (global or node-level) + MergedServices string // Merged services configuration from all imports + MergedNetwork string // Merged network configuration from all imports + MergedPermissions string // Merged permissions configuration from all imports + MergedSecretMasking string // Merged secret-masking steps from all imports + MergedBots []string // Merged bots list from all imports (union of bot names) + MergedSkipRoles []string // Merged skip-roles list from all imports (union of role names) + MergedSkipBots []string // Merged skip-bots list from all imports (union of usernames) + MergedActivationGitHubToken string // GitHub token from on.github-token in first imported workflow that defines it + MergedActivationGitHubApp string // JSON-encoded on.github-app from first imported workflow that defines it + MergedTopLevelGitHubApp string // JSON-encoded top-level github-app from first imported workflow that defines it + MergedCheckout string // JSON-encoded checkout configurations from imported workflows (one JSON value per line) + MergedPostSteps string // Merged post-steps configuration from all imports (appended in order) + MergedLabels []string // Merged labels from all imports (union of label names) + MergedCaches []string // Merged cache configurations from all imports (appended in order) + MergedJobs string // Merged jobs from imported YAML workflows (JSON format) + MergedEnv string // Merged env configuration from all imports (JSON format) + MergedEnvSources map[string]string // env var name → source import path (for conflict detection and lock file header listing) + MergedFeatures []map[string]any // Merged features configuration from all imports (parsed YAML structures) + MergedModels []map[string][]string // Merged model alias definitions from all imports (first import to define a key wins among imports) + MergedObservability string // Observability config (JSON) from first import that defines it (first-wins) + ImportedFiles []string // List of imported file paths (for manifest) + AgentFile string // Path to custom agent file (if imported) + AgentImportSpec string // Original import specification for agent file (e.g., "owner/repo/path@ref") + RepositoryImports []string // List of repository imports (format: "owner/repo@ref") for .github folder merging // ImportInputs uses map[string]any because input values can be different types (string, number, boolean). // This is parsed from YAML frontmatter where the structure is dynamic and not known at compile time. // This is an appropriate use of 'any' for dynamic YAML/JSON data. diff --git a/pkg/parser/schemas/main_workflow_schema.json b/pkg/parser/schemas/main_workflow_schema.json index 9f5a74f97aa..df0bbfee5a5 100644 --- a/pkg/parser/schemas/main_workflow_schema.json +++ b/pkg/parser/schemas/main_workflow_schema.json @@ -2730,6 +2730,24 @@ } ] }, + "models": { + "description": "Named model alias definitions with ordered fallback lists, resolved recursively by AWF. Each key is an alias name (use empty string \"\" for the default policy). Each value is an ordered list of vendor/modelid glob patterns or other alias names to try in sequence. Entries defined here are merged on top of the builtin aliases; the main workflow file always wins over imported aliases. Builtin aliases include: sonnet, haiku, opus, gpt-5, gpt-5-mini, gpt-5-codex, gemini-flash, gemini-pro, small, mini, large, auto.", + "type": "object", + "additionalProperties": { + "type": "array", + "items": { + "type": "string", + "description": "A vendor/modelid glob pattern (e.g. 'copilot/*sonnet*') or an alias name (e.g. 'sonnet') to resolve recursively." + }, + "description": "Ordered list of model patterns or alias names to try in sequence." + }, + "examples": [ + { + "sonnet": ["mygateway/*sonnet-v3*"], + "": ["sonnet", "gpt-5-codex"] + } + ] + }, "experiments": { "description": "A/B testing experiments. Each key is an experiment name; the value is either an array of two or more variant strings (bare-array form) or an object with a 'variants' field plus optional metadata fields (description, metric, weight, issue, start_date, end_date, hypothesis, secondary_metrics, guardrail_metrics, min_samples, owner). At runtime the activation job picks a variant using actions/cache to maintain consistent assignment across runs. Use ${{ experiments. }} in the workflow prompt to reference the selected variant. When multiple experiments are declared, assignments are statistically balanced using a counter that round-robins across variants (or weighted when 'weight' is provided).", "type": "object", diff --git a/pkg/workflow/awf_config.go b/pkg/workflow/awf_config.go index 1094a457a98..be85815e9e2 100644 --- a/pkg/workflow/awf_config.go +++ b/pkg/workflow/awf_config.go @@ -68,6 +68,17 @@ type AWFConfigFile struct { // Container contains container execution configuration. Container *AWFContainerConfig `json:"container,omitempty"` + + // Models contains model alias and fallback policy definitions. + // Keys are alias names (empty string "" = default policy); values are ordered + // lists of vendor/modelid patterns or other alias names to try in sequence. + // AWF resolves aliases recursively; loops are not permitted. + // + // NOTE: Pending AWF binary support (config.models is not yet recognised by the + // AWF firewall schema). This field is intentionally omitted from JSON output + // until the AWF schema at awf-config.v1.json is updated to include "models". + // The field remains here so the struct is ready once AWF support lands. + Models map[string][]string `json:"-"` } // AWFNetworkConfig is the "network" section of the AWF config file. @@ -209,6 +220,12 @@ func BuildAWFConfigJSON(config AWFCommandConfig) (string, error) { awfConfigLog.Printf("Container section: image_tag=%s", awfImageTag) } + // ── Models section ──────────────────────────────────────────────────────── + if config.WorkflowData != nil && len(config.WorkflowData.ModelMappings) > 0 { + awfConfig.Models = config.WorkflowData.ModelMappings + awfConfigLog.Printf("Models section: %d alias entries", len(config.WorkflowData.ModelMappings)) + } + jsonBytes, err := json.Marshal(awfConfig) if err != nil { return "", fmt.Errorf("failed to marshal AWF config to JSON: %w", err) diff --git a/pkg/workflow/compiler_types.go b/pkg/workflow/compiler_types.go index 636c3e5fe41..28f19369475 100644 --- a/pkg/workflow/compiler_types.go +++ b/pkg/workflow/compiler_types.go @@ -519,6 +519,7 @@ type WorkflowData struct { CachedAllowedDomainsStr string // cached allowed-domains string for sanitization (for performance optimization); computed once and reused across multiple compilation steps CachedAllowedDomainsComputed bool // true once CachedAllowedDomainsStr has been set; distinguishes "computed empty" from "not yet computed" KnownActionCredentialEnvVars map[string]bool // env vars for clean_known_action_credentials.sh; keyed by GH_AW_CLEAN_* names; nil when no known credential-leaking actions are detected + ModelMappings map[string][]string // merged model alias map (builtins + imported workflow aliases + main frontmatter overrides, in priority order); NOT yet emitted to AWF config JSON — pending AWF firewall support (config.models) } // PinContext returns an actionpins.PinContext backed by this WorkflowData. diff --git a/pkg/workflow/frontmatter_types.go b/pkg/workflow/frontmatter_types.go index 7859142f0fb..47cd94f71a7 100644 --- a/pkg/workflow/frontmatter_types.go +++ b/pkg/workflow/frontmatter_types.go @@ -275,6 +275,12 @@ type FrontmatterConfig struct { // Experiments during frontmatter parsing. Keys match those of Experiments. ExperimentConfigs map[string]*ExperimentConfig `json:"-"` + // Model aliases and fallback policies. + // Keys are alias names (empty string "" = default policy); values are ordered lists of + // model patterns or alias references to try in sequence. + // Merged with the builtin model aliases at compile time; frontmatter entries take precedence. + Models map[string][]string `json:"models,omitempty"` + // Rate limiting configuration RateLimit *RateLimitConfig `json:"rate-limit,omitempty"` diff --git a/pkg/workflow/model_aliases.go b/pkg/workflow/model_aliases.go new file mode 100644 index 00000000000..7fc8d65502b --- /dev/null +++ b/pkg/workflow/model_aliases.go @@ -0,0 +1,149 @@ +// This file provides model alias and fallback resolution for AWF (Agentic Workflow Firewall). +// +// # Model Alias Format +// +// A model payload is a map from alias name to an ordered list of model patterns: +// +// { +// "sonnet": ["copilot/*sonnet*", "anthropic/*sonnet*"], +// "haiku": ["copilot/*haiku*", "anthropic/*haiku*"], +// "": ["sonnet", "gpt-5"] // default policy +// } +// +// The syntax for each pattern entry is: +// - "vendor/modelid" — exact vendor-scoped model name +// - "vendor/model*id" — wildcard pattern (supports * as a glob wildcard) +// - "alias" — reference to another alias in the same map (recursive resolution) +// +// AWF resolves aliases recursively. Loops are not permitted. +// +// # Builtin Aliases +// +// gh-aw ships a set of builtin model aliases that cover the major model families. +// Frontmatter-defined aliases are merged on top of the builtins, allowing workflows +// to extend or override the defaults without replacing the entire mapping. + +package workflow + +import "maps" + +// BuiltinModelAliases returns the built-in model alias map that covers the main +// model families supported by gh-aw. The returned map is a freshly allocated +// copy so callers may freely modify it. +// +// Vendor aliases (patterns use * as a glob wildcard, prefer copilot gateway first): +// - "sonnet" → Anthropic Sonnet family +// - "haiku" → Anthropic Haiku family +// - "opus" → Anthropic Opus family +// - "gpt-5" → OpenAI GPT-5 family +// - "gpt-5-mini" → OpenAI GPT-5-mini family +// - "gpt-5-codex" → OpenAI GPT-5-Codex family +// - "gemini-flash" → Google Gemini Flash family (fast/lightweight) +// - "gemini-pro" → Google Gemini Pro family (full-capability) +// +// Meta-aliases (reference other aliases; resolved recursively by AWF): +// - "mini" → haiku, gpt-5-mini, gemini-flash +// - "large" → sonnet, gpt-5, gemini-pro +// - "auto" → large (convenience alias for the default capable tier) +func BuiltinModelAliases() map[string][]string { + return map[string][]string{ + // ── Anthropic ──────────────────────────────────────────────────────── + "sonnet": { + "copilot/*sonnet*", + "anthropic/*sonnet*", + }, + "haiku": { + "copilot/*haiku*", + "anthropic/*haiku*", + }, + "opus": { + "copilot/*opus*", + "anthropic/*opus*", + }, + // ── OpenAI ─────────────────────────────────────────────────────────── + "gpt-5": { + "copilot/gpt-5*", + "openai/gpt-5*", + }, + "gpt-5-mini": { + "copilot/gpt-5*mini*", + "openai/gpt-5*mini*", + }, + "gpt-5-codex": { + "copilot/gpt-5*codex*", + "openai/gpt-5*codex*", + }, + // ── Google ─────────────────────────────────────────────────────────── + "gemini-flash": { + "copilot/gemini-*flash*", + "google/gemini-*flash*", + }, + "gemini-pro": { + "copilot/gemini-*pro*", + "google/gemini-*pro*", + }, + // ── Meta-aliases ───────────────────────────────────────────────────── + // These reference other aliases; AWF resolves them recursively. + // "small" — same as "mini" (convenience alias for lightweight/fast models). + // "mini" — lightweight/fast models across all supported vendors. + // "large" — full-capability models across all supported vendors. + // "auto" — convenience alias that resolves to the "large" tier. + "small": { + "mini", + }, + "mini": { + "haiku", + "gpt-5-mini", + "gemini-flash", + }, + "large": { + "sonnet", + "gpt-5", + "gemini-pro", + }, + "auto": { + "large", + }, + } +} + +// MergeImportedModelAliases builds the final model alias map from three layers, +// with later layers overriding earlier ones (highest priority last): +// +// 1. Builtin aliases (lowest priority) +// 2. Imported workflow aliases — merged in import order; first import to define a +// key wins among imports (same "first-wins among peers" semantics as features). +// 3. Main workflow frontmatter aliases (highest priority — main workflow file wins) +// +// If both importedModels and frontmatterModels are nil/empty, the builtin aliases are +// returned as-is (identical to MergeModelAliases(nil)). +func MergeImportedModelAliases(importedModels []map[string][]string, frontmatterModels map[string][]string) map[string][]string { + merged := BuiltinModelAliases() + + // Layer 2 — imported models (first import to define a key wins among imports). + for _, importedMap := range importedModels { + for k, v := range importedMap { + if _, exists := merged[k]; !exists { + merged[k] = v + } + } + } + + // Layer 3 — main workflow frontmatter always wins. + maps.Copy(merged, frontmatterModels) + + return merged +} + +// MergeModelAliases merges the frontmatter-defined model aliases on top of the +// builtin aliases and returns the combined map. Frontmatter entries always take +// precedence: if the same key exists in both the builtins and the frontmatter +// definition, the frontmatter value replaces the builtin value entirely. +// +// If frontmatterModels is nil or empty, the builtin aliases are returned as-is. +// +// For the full three-layer merge that also incorporates imported workflow aliases, +// use MergeImportedModelAliases. +func MergeModelAliases(frontmatterModels map[string][]string) map[string][]string { + return MergeImportedModelAliases(nil, frontmatterModels) +} diff --git a/pkg/workflow/model_aliases_import_test.go b/pkg/workflow/model_aliases_import_test.go new file mode 100644 index 00000000000..3de16872934 --- /dev/null +++ b/pkg/workflow/model_aliases_import_test.go @@ -0,0 +1,158 @@ +//go:build !integration + +package workflow_test + +import ( + "os" + "path/filepath" + "testing" + + "github.com/github/gh-aw/pkg/stringutil" + "github.com/github/gh-aw/pkg/testutil" + "github.com/github/gh-aw/pkg/workflow" + + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" +) + +// TestModelAliasesFromImportedWorkflow verifies that model aliases defined in an imported +// (shared) workflow file are included in the compiled AWF config, and that the main +// workflow file's aliases take precedence. +func TestModelAliasesFromImportedWorkflow(t *testing.T) { + tempDir := testutil.TempDir(t, "models-import-test-*") + + // Shared workflow defines a custom alias and a non-conflicting vendor alias. + sharedWorkflowPath := filepath.Join(tempDir, "shared-models.md") + sharedWorkflowContent := `--- +models: + shared-alias: + - shared/model-v1 + sonnet: + - import/sonnet-override +--- + +# Shared Models +` + require.NoError(t, os.WriteFile(sharedWorkflowPath, []byte(sharedWorkflowContent), 0644)) + + // Main workflow imports the shared file and overrides one alias. + mainWorkflowPath := filepath.Join(tempDir, "main-workflow.md") + mainWorkflowContent := `--- +on: issues +permissions: + contents: read + issues: read +engine: copilot +models: + main-alias: + - main/model-v1 + shared-alias: + - main/override-shared +imports: + - shared-models.md +--- + +# Main Workflow +` + require.NoError(t, os.WriteFile(mainWorkflowPath, []byte(mainWorkflowContent), 0644)) + + compiler := workflow.NewCompiler() + require.NoError(t, compiler.CompileWorkflow(mainWorkflowPath), "workflow compilation should succeed") + + lockFilePath := stringutil.MarkdownToLockFile(mainWorkflowPath) + lockFileContent, err := os.ReadFile(lockFilePath) + require.NoError(t, err, "lock file should be readable") + + lockYAML := string(lockFileContent) + + // Verify the generated lock file compiles successfully. Since models are not yet + // emitted to awf-config.json (pending AWF firewall support), we verify that the + // compilation itself succeeds and produces a valid lock file without checking for + // model alias names in the lock YAML (they don't appear in the JSON config yet). + assert.NotEmpty(t, lockYAML, "lock file should be non-empty after successful compilation") +} + +// TestModelAliasesImportMergeOrder verifies the priority order: +// builtins < imported aliases < main workflow aliases. +func TestModelAliasesImportMergeOrder(t *testing.T) { + t.Run("imported alias visible, main wins over import", func(t *testing.T) { + imported := []map[string][]string{ + { + "import-alias": {"import/model"}, + "shared-key": {"import/shared"}, + }, + } + mainModels := map[string][]string{ + "main-alias": {"main/model"}, + "shared-key": {"main/shared"}, + } + + merged := workflow.MergeImportedModelAliases(imported, mainModels) + + // Imported alias that main doesn't touch is visible. + assert.Equal(t, []string{"import/model"}, merged["import-alias"], + "import-alias from imported workflow should be in the merged map") + + // Main alias beats import for the same key. + assert.Equal(t, []string{"main/shared"}, merged["shared-key"], + "main workflow alias should win over imported alias for same key") + + // Main-only alias is present. + assert.Equal(t, []string{"main/model"}, merged["main-alias"], + "main-only alias should be present") + + // Builtins are still present. + assert.NotEmpty(t, merged["sonnet"], "builtin sonnet should still be present") + assert.NotEmpty(t, merged["auto"], "builtin auto should still be present") + }) +} + +// TestModelAliasesAWFConfigJSON verifies that model alias entries from imported workflows +// are merged into WorkflowData.ModelMappings during compilation. +// +// NOTE: The "models" field is intentionally excluded from the AWF config JSON until the +// AWF firewall binary is updated to recognise config.models. Assertions check +// ModelMappings directly rather than the serialised JSON. +func TestModelAliasesAWFConfigJSON(t *testing.T) { + awfConfig := workflow.AWFCommandConfig{ + EngineName: "copilot", + AllowedDomains: "github.com", + WorkflowData: &workflow.WorkflowData{ + EngineConfig: &workflow.EngineConfig{ID: "copilot"}, + NetworkPermissions: &workflow.NetworkPermissions{ + Firewall: &workflow.FirewallConfig{Enabled: true}, + }, + // Simulate: import defines a new alias, main overrides a builtin. + ModelMappings: workflow.MergeImportedModelAliases( + []map[string][]string{ + {"import-alias": {"import/model"}}, + }, + map[string][]string{ + "haiku": {"main/haiku-override"}, + }, + ), + }, + } + + jsonStr, err := workflow.BuildAWFConfigJSON(awfConfig) + require.NoError(t, err, "BuildAWFConfigJSON should not return an error") + + // models must NOT appear in the JSON until the AWF binary supports it + assert.NotContains(t, jsonStr, `"models"`, "models section must be absent from AWF config JSON until AWF binary supports it") + + // Verify that the alias map is correctly populated in WorkflowData. + mappings := awfConfig.WorkflowData.ModelMappings + require.NotNil(t, mappings, "ModelMappings should be set on WorkflowData") + + // Imported alias is in the model mappings. + assert.Equal(t, []string{"import/model"}, mappings["import-alias"], + "import-alias from imported workflow should be in ModelMappings") + + // Main workflow override wins over builtin haiku. + assert.Equal(t, []string{"main/haiku-override"}, mappings["haiku"], + "main workflow alias should override builtin haiku in ModelMappings") + + // Other builtins preserved. + assert.NotEmpty(t, mappings["sonnet"], "builtin sonnet should still be in ModelMappings") + assert.NotEmpty(t, mappings["auto"], "builtin auto should still be in ModelMappings") +} diff --git a/pkg/workflow/model_aliases_test.go b/pkg/workflow/model_aliases_test.go new file mode 100644 index 00000000000..16e7942fe47 --- /dev/null +++ b/pkg/workflow/model_aliases_test.go @@ -0,0 +1,320 @@ +//go:build !integration + +package workflow + +import ( + "encoding/json" + "testing" + + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" +) + +// TestBuiltinModelAliases verifies that the builtin model alias map covers the main +// model families and returns a fresh map on each call. +func TestBuiltinModelAliases(t *testing.T) { + aliases := BuiltinModelAliases() + + expectedFamilies := []string{ + "sonnet", "haiku", "opus", + "gpt-5", "gpt-5-mini", "gpt-5-codex", + "gemini-flash", "gemini-pro", + "mini", "large", "auto", + } + for _, family := range expectedFamilies { + patterns, ok := aliases[family] + assert.True(t, ok, "expected builtin alias for family %q", family) + assert.NotEmpty(t, patterns, "builtin alias %q should have at least one pattern", family) + } + + // Vendor aliases should include at least one copilot/* pattern. + // Meta-aliases (mini, large, auto) reference other alias names and are excluded here. + vendorFamilies := []string{"sonnet", "haiku", "opus", "gpt-5", "gpt-5-mini", "gpt-5-codex", "gemini-flash", "gemini-pro"} + for _, family := range vendorFamilies { + patterns := aliases[family] + hasCopilot := false + for _, p := range patterns { + if len(p) > 7 && p[:7] == "copilot" { + hasCopilot = true + break + } + } + assert.True(t, hasCopilot, "builtin alias %q should include a copilot/* pattern", family) + } + + // Meta-aliases reference other alias names (resolved recursively by AWF). + assert.Equal(t, []string{"haiku", "gpt-5-mini", "gemini-flash"}, aliases["mini"], "mini should reference haiku, gpt-5-mini, and gemini-flash") + assert.Equal(t, []string{"sonnet", "gpt-5", "gemini-pro"}, aliases["large"], "large should reference sonnet, gpt-5, and gemini-pro") + assert.Equal(t, []string{"large"}, aliases["auto"], "auto should fall back to large") + + // Returns a fresh copy — mutating one call's map must not affect another call. + aliases["sonnet"] = []string{"custom/model"} + aliases2 := BuiltinModelAliases() + assert.NotEqual(t, aliases["sonnet"], aliases2["sonnet"], "BuiltinModelAliases should return a fresh copy each time") +} + +// TestMergeModelAliases verifies that frontmatter-defined aliases are merged on top +// of the builtins. +func TestMergeModelAliases(t *testing.T) { + t.Run("nil frontmatter returns all builtins", func(t *testing.T) { + merged := MergeModelAliases(nil) + builtins := BuiltinModelAliases() + assert.Len(t, merged, len(builtins), "nil frontmatter should return exactly the builtins") + for k, v := range builtins { + assert.Equal(t, v, merged[k], "builtin alias %q should be present unchanged", k) + } + }) + + t.Run("empty frontmatter returns all builtins", func(t *testing.T) { + merged := MergeModelAliases(map[string][]string{}) + builtins := BuiltinModelAliases() + assert.Len(t, merged, len(builtins), "empty frontmatter should return exactly the builtins") + }) + + t.Run("frontmatter override replaces builtin entry", func(t *testing.T) { + custom := map[string][]string{ + "sonnet": {"myvendor/sonnet-custom"}, + } + merged := MergeModelAliases(custom) + assert.Equal(t, []string{"myvendor/sonnet-custom"}, merged["sonnet"], + "frontmatter override should replace the builtin sonnet alias") + // Other builtins should be unaffected. + assert.NotEmpty(t, merged["haiku"], "haiku builtin should still be present") + }) + + t.Run("frontmatter adds new alias", func(t *testing.T) { + custom := map[string][]string{ + "my-alias": {"copilot/my-model"}, + } + merged := MergeModelAliases(custom) + assert.Equal(t, []string{"copilot/my-model"}, merged["my-alias"], + "new frontmatter alias should be present in merged map") + // Builtins should still be present. + assert.NotEmpty(t, merged["sonnet"], "sonnet builtin should still be present") + }) + + t.Run("default policy key is supported", func(t *testing.T) { + custom := map[string][]string{ + "": {"sonnet", "gpt-5-codex"}, + } + merged := MergeModelAliases(custom) + assert.Equal(t, []string{"sonnet", "gpt-5-codex"}, merged[""], + "default policy (empty key) should be stored and returned") + }) +} + +// TestBuildAWFConfigJSON_ModelsSection verifies model alias behaviour in BuildAWFConfigJSON. +// +// NOTE: The "models" field is intentionally excluded from the AWF config JSON until the +// AWF firewall binary is updated to recognise config.models (awf-config.v1.json schema). +// The model alias infrastructure (builtin aliases, frontmatter overrides, import merging) +// remains fully operational inside gh-aw; once AWF support lands the json:"-" tag on +// AWFConfigFile.Models can be changed to json:"models,omitempty" to re-enable emission. +func TestBuildAWFConfigJSON_ModelsSection(t *testing.T) { + t.Run("builtin model aliases are included when WorkflowData has ModelMappings", func(t *testing.T) { + config := AWFCommandConfig{ + EngineName: "copilot", + AllowedDomains: "github.com", + WorkflowData: &WorkflowData{ + EngineConfig: &EngineConfig{ID: "copilot"}, + NetworkPermissions: &NetworkPermissions{ + Firewall: &FirewallConfig{Enabled: true}, + }, + ModelMappings: MergeModelAliases(nil), + }, + } + + jsonStr, err := BuildAWFConfigJSON(config) + require.NoError(t, err, "BuildAWFConfigJSON should not return an error") + + var parsed map[string]any + require.NoError(t, json.Unmarshal([]byte(jsonStr), &parsed), "result must be valid JSON") + + // models must NOT appear in the JSON until the AWF binary supports it + assert.NotContains(t, parsed, "models", "models section must be absent from AWF config JSON until AWF binary supports it") + + // but the alias map is still populated in WorkflowData + assert.NotEmpty(t, config.WorkflowData.ModelMappings, "ModelMappings should be populated on WorkflowData") + assert.Contains(t, config.WorkflowData.ModelMappings, "sonnet", "ModelMappings should include sonnet alias") + assert.Contains(t, config.WorkflowData.ModelMappings, "haiku", "ModelMappings should include haiku alias") + assert.Contains(t, config.WorkflowData.ModelMappings, "auto", "ModelMappings should include auto alias") + }) + + t.Run("frontmatter override is reflected in WorkflowData but not in AWF config JSON", func(t *testing.T) { + custom := map[string][]string{ + "sonnet": {"myvendor/sonnet-v3"}, + "": {"sonnet"}, + } + config := AWFCommandConfig{ + EngineName: "copilot", + AllowedDomains: "github.com", + WorkflowData: &WorkflowData{ + EngineConfig: &EngineConfig{ID: "copilot"}, + NetworkPermissions: &NetworkPermissions{ + Firewall: &FirewallConfig{Enabled: true}, + }, + ModelMappings: MergeModelAliases(custom), + }, + } + + jsonStr, err := BuildAWFConfigJSON(config) + require.NoError(t, err, "BuildAWFConfigJSON should not return an error") + + // models must NOT appear in the JSON until the AWF binary supports it + assert.NotContains(t, jsonStr, `"models"`, "models section must be absent from AWF config JSON until AWF binary supports it") + + // but frontmatter overrides are visible in WorkflowData + assert.Equal(t, []string{"myvendor/sonnet-v3"}, config.WorkflowData.ModelMappings["sonnet"], + "frontmatter override for sonnet should be stored in ModelMappings") + assert.Equal(t, []string{"sonnet"}, config.WorkflowData.ModelMappings[""], + "default policy should be stored in ModelMappings") + }) + + t.Run("no models section when ModelMappings is nil", func(t *testing.T) { + config := AWFCommandConfig{ + EngineName: "copilot", + AllowedDomains: "github.com", + WorkflowData: &WorkflowData{ + EngineConfig: &EngineConfig{ID: "copilot"}, + NetworkPermissions: &NetworkPermissions{ + Firewall: &FirewallConfig{Enabled: true}, + }, + ModelMappings: nil, + }, + } + + jsonStr, err := BuildAWFConfigJSON(config) + require.NoError(t, err, "BuildAWFConfigJSON should not return an error") + + assert.NotContains(t, jsonStr, `"models"`, "models section should be absent when ModelMappings is nil") + }) +} + +// TestMergeImportedModelAliases verifies the three-layer merge: builtins → imports → main. +func TestMergeImportedModelAliases(t *testing.T) { + t.Run("no imports and no frontmatter returns builtins", func(t *testing.T) { + merged := MergeImportedModelAliases(nil, nil) + builtins := BuiltinModelAliases() + assert.Len(t, merged, len(builtins), "should return exactly the builtins") + for k, v := range builtins { + assert.Equal(t, v, merged[k], "builtin alias %q should be present unchanged", k) + } + }) + + t.Run("imported alias is added when not in builtins", func(t *testing.T) { + imported := []map[string][]string{ + {"my-imported": {"vendor/imported-model"}}, + } + merged := MergeImportedModelAliases(imported, nil) + assert.Equal(t, []string{"vendor/imported-model"}, merged["my-imported"], + "imported alias should be present in merged map") + assert.NotEmpty(t, merged["sonnet"], "builtin sonnet should still be present") + }) + + t.Run("import cannot override a builtin alias", func(t *testing.T) { + imported := []map[string][]string{ + {"sonnet": {"imported/sonnet-override"}}, + } + merged := MergeImportedModelAliases(imported, nil) + builtins := BuiltinModelAliases() + assert.Equal(t, builtins["sonnet"], merged["sonnet"], + "import should NOT override a builtin alias; builtin takes precedence over import") + }) + + t.Run("first import wins among multiple imports for the same key", func(t *testing.T) { + imported := []map[string][]string{ + {"shared-alias": {"first-import/model"}}, + {"shared-alias": {"second-import/model"}}, + } + merged := MergeImportedModelAliases(imported, nil) + assert.Equal(t, []string{"first-import/model"}, merged["shared-alias"], + "first import should win among competing imports for the same alias key") + }) + + t.Run("main workflow frontmatter overrides imported alias", func(t *testing.T) { + imported := []map[string][]string{ + {"my-alias": {"import/model"}}, + } + frontmatter := map[string][]string{ + "my-alias": {"main/model"}, + } + merged := MergeImportedModelAliases(imported, frontmatter) + assert.Equal(t, []string{"main/model"}, merged["my-alias"], + "main workflow frontmatter should win over imported alias") + }) + + t.Run("main workflow frontmatter overrides builtin alias", func(t *testing.T) { + frontmatter := map[string][]string{ + "sonnet": {"mygateway/sonnet-v3"}, + } + merged := MergeImportedModelAliases(nil, frontmatter) + assert.Equal(t, []string{"mygateway/sonnet-v3"}, merged["sonnet"], + "main workflow frontmatter should override builtin sonnet alias") + assert.NotEmpty(t, merged["haiku"], "other builtins should still be present") + }) + + t.Run("all three layers are combined correctly", func(t *testing.T) { + imported := []map[string][]string{ + { + "import-only": {"import/model"}, + "both": {"import/both"}, + "sonnet": {"import/sonnet"}, // shadowed by builtin + }, + } + frontmatter := map[string][]string{ + "main-only": {"main/model"}, + "both": {"main/both"}, + } + merged := MergeImportedModelAliases(imported, frontmatter) + + // import-only key comes from import (no conflict) + assert.Equal(t, []string{"import/model"}, merged["import-only"], + "import-only alias should come from the import layer") + + // main-only key comes from main workflow + assert.Equal(t, []string{"main/model"}, merged["main-only"], + "main-only alias should come from the main workflow layer") + + // 'both' key: main workflow wins over import + assert.Equal(t, []string{"main/both"}, merged["both"], + "main workflow should win over import for the 'both' key") + + // 'sonnet' key: builtin wins over import + builtins := BuiltinModelAliases() + assert.Equal(t, builtins["sonnet"], merged["sonnet"], + "builtin should win over import for the 'sonnet' key") + }) +} + +// correctly by ParseFrontmatterConfig. +func TestFrontmatterModelsField(t *testing.T) { + t.Run("models field is parsed from frontmatter", func(t *testing.T) { + frontmatter := map[string]any{ + "name": "test-workflow", + "models": map[string]any{ + "my-model": []any{"copilot/my-model-v1", "openai/my-model-v1"}, + "": []any{"my-model"}, + }, + } + + config, err := ParseFrontmatterConfig(frontmatter) + require.NoError(t, err, "ParseFrontmatterConfig should succeed with models field") + require.NotNil(t, config, "parsed config should not be nil") + + assert.Equal(t, []string{"copilot/my-model-v1", "openai/my-model-v1"}, config.Models["my-model"], + "models[my-model] should be parsed correctly") + assert.Equal(t, []string{"my-model"}, config.Models[""], + "models default policy (empty key) should be parsed correctly") + }) + + t.Run("models field is optional", func(t *testing.T) { + frontmatter := map[string]any{ + "name": "test-workflow", + } + + config, err := ParseFrontmatterConfig(frontmatter) + require.NoError(t, err, "ParseFrontmatterConfig should succeed without models field") + require.NotNil(t, config, "parsed config should not be nil") + assert.Nil(t, config.Models, "models should be nil when not specified in frontmatter") + }) +} diff --git a/pkg/workflow/workflow_builder.go b/pkg/workflow/workflow_builder.go index 6e7349e9204..c076055a1bd 100644 --- a/pkg/workflow/workflow_builder.go +++ b/pkg/workflow/workflow_builder.go @@ -132,6 +132,15 @@ func (c *Compiler) buildInitialWorkflowData( } } + // Populate model mappings: merge builtin aliases, any imported-workflow aliases, and + // main-workflow frontmatter overrides. Priority (highest last): + // builtins → imported workflow aliases → main workflow frontmatter (main wins). + var frontmatterModels map[string][]string + if toolsResult.parsedFrontmatter != nil { + frontmatterModels = toolsResult.parsedFrontmatter.Models + } + workflowData.ModelMappings = MergeImportedModelAliases(importsResult.MergedModels, frontmatterModels) + return workflowData }