diff --git a/pkg/workflow/behavior_defined_engine.go b/pkg/workflow/behavior_defined_engine.go index eaf2580b507..5dd16e5fef9 100644 --- a/pkg/workflow/behavior_defined_engine.go +++ b/pkg/workflow/behavior_defined_engine.go @@ -7,6 +7,7 @@ import ( "maps" "slices" "strings" + "sync" "time" "github.com/github/gh-aw/pkg/constants" @@ -502,10 +503,48 @@ func isEngineDefinitionForm(def *EngineDefinition) bool { return def.Models.Default != "" || len(def.Models.Supported) > 0 || len(def.Options) > 0 } +// engineDefinitionBuiltinKeys is the set of JSON strings corresponding to +// built-in engine definitions. It is populated once at startup by +// loadBuiltinEngineDefinitions (via registerBuiltinEngineDefinitionJSON) and +// never modified afterward. Only JSON keys present in this set are eligible for +// caching in engineDefinitionCache, which prevents unbounded growth in long-lived +// compile --watch sessions where each edit to a custom engine definition would +// otherwise create a new cache entry. +var engineDefinitionBuiltinKeys sync.Map // map[string]struct{} + +// registerBuiltinEngineDefinitionJSON marks jsonKey as a known built-in engine +// JSON string so that parseEngineDefinitionFromJSON will cache it. +func registerBuiltinEngineDefinitionJSON(jsonKey string) { + engineDefinitionBuiltinKeys.Store(jsonKey, struct{}{}) +} + +// engineDefinitionCache caches parsed EngineDefinition values for built-in +// engine JSON strings. Built-in engine files are injected as imports on every +// CompileWorkflow call and their JSON representation is always identical, so +// caching avoids the repeated JSON→any→YAML→struct round-trip that accounted +// for ~24% of BenchmarkCompileMCPWorkflow wall-clock time. +// +// Only keys present in engineDefinitionBuiltinKeys are stored to bound the cache +// to the fixed set of built-in engines. Deep copies are returned on every lookup +// so callers cannot corrupt the cached state through mutations to pointers, slices, +// or maps within the returned definition. +var engineDefinitionCache sync.Map // map[string]EngineDefinition + func parseEngineDefinitionFromJSON(engineJSON string) (*EngineDefinition, error) { if engineJSON == "" { return nil, nil } + // Fast path: return a deep copy of the cached definition when available. + if cached, ok := engineDefinitionCache.Load(engineJSON); ok { + if def, ok := cached.(EngineDefinition); ok { + defCopy := deepCopyEngineDefinition(def) + return &defCopy, nil + } + // Type assertion failure indicates cache corruption or a concurrent Store with + // an unexpected type. Log and fall through to re-parse so the caller still works. + behaviorDefinedEngineLog.Printf("engineDefinitionCache: unexpected value type %T for key (len=%d); re-parsing", cached, len(engineJSON)) + engineDefinitionCache.Delete(engineJSON) + } var engineData any if err := json.Unmarshal([]byte(engineJSON), &engineData); err != nil { return nil, fmt.Errorf("failed to parse engine JSON: %w", err) @@ -524,9 +563,153 @@ func parseEngineDefinitionFromJSON(engineJSON string) (*EngineDefinition, error) if def.RuntimeID == "" { def.RuntimeID = def.ID } + // Cache only built-in engine definitions (keys pre-seeded by loadBuiltinEngineDefinitions) + // to prevent unbounded memory growth. Store a deep copy so that any mutations the + // caller makes to the returned definition cannot corrupt the cached entry. + if _, isBuiltin := engineDefinitionBuiltinKeys.Load(engineJSON); isBuiltin { + cacheCopy := deepCopyEngineDefinition(def) + engineDefinitionCache.Store(engineJSON, cacheCopy) + } return &def, nil } +// deepCopyAny returns a fully independent copy of v for values produced by +// yaml.Unmarshal into interface{}. The possible concrete types are: +// nil, bool, int, float64, string, []any, and map[string]any. +func deepCopyAny(v any) any { + switch val := v.(type) { + case map[string]any: + cp := make(map[string]any, len(val)) + for k, elem := range val { + cp[k] = deepCopyAny(elem) + } + return cp + case []any: + cp := make([]any, len(val)) + for i, elem := range val { + cp[i] = deepCopyAny(elem) + } + return cp + default: + // Scalars (nil, bool, int, float64, string) are immutable value types. + return v + } +} + +// deepCopyEngineDefinition returns a fully independent copy of src. All reference +// types (pointers, slices, maps) are recursively copied so that neither the caller +// nor the cache can corrupt the other's state through shared references. +func deepCopyEngineDefinition(src EngineDefinition) EngineDefinition { + dst := src // value copy covers all scalar fields + + // Models.Supported + if src.Models.Supported != nil { + dst.Models.Supported = make([]string, len(src.Models.Supported)) + copy(dst.Models.Supported, src.Models.Supported) + } + + // Auth ([]AuthBinding; elements contain only string fields so element copy suffices) + if src.Auth != nil { + dst.Auth = make([]AuthBinding, len(src.Auth)) + copy(dst.Auth, src.Auth) + } + + // Options (map[string]any; values may contain nested maps or slices from YAML unmarshal) + if src.Options != nil { + dst.Options = make(map[string]any, len(src.Options)) + for k, v := range src.Options { + dst.Options[k] = deepCopyAny(v) + } + } + + // Provider.Auth + if src.Provider.Auth != nil { + authCopy := *src.Provider.Auth // AuthDefinition contains only string fields + dst.Provider.Auth = &authCopy + } + + // Provider.Request + if src.Provider.Request != nil { + reqCopy := *src.Provider.Request + if src.Provider.Request.Query != nil { + reqCopy.Query = make(map[string]string, len(src.Provider.Request.Query)) + maps.Copy(reqCopy.Query, src.Provider.Request.Query) + } + if src.Provider.Request.BodyInject != nil { + reqCopy.BodyInject = make(map[string]string, len(src.Provider.Request.BodyInject)) + maps.Copy(reqCopy.BodyInject, src.Provider.Request.BodyInject) + } + dst.Provider.Request = &reqCopy + } + + // Behaviors + if src.Behaviors != nil { + behaviorsCopy := deepCopyEngineBehaviorDefinition(*src.Behaviors) + dst.Behaviors = &behaviorsCopy + } + + return dst +} + +// deepCopyEngineBehaviorDefinition returns a fully independent copy of src. +func deepCopyEngineBehaviorDefinition(src EngineBehaviorDefinition) EngineBehaviorDefinition { + dst := src // value copy covers all scalar fields + + // SupportedEnvVarKeys + if src.SupportedEnvVarKeys != nil { + dst.SupportedEnvVarKeys = make([]string, len(src.SupportedEnvVarKeys)) + copy(dst.SupportedEnvVarKeys, src.SupportedEnvVarKeys) + } + + // Manifest + if src.Manifest != nil { + manifestCopy := *src.Manifest + if src.Manifest.Files != nil { + manifestCopy.Files = make([]string, len(src.Manifest.Files)) + copy(manifestCopy.Files, src.Manifest.Files) + } + if src.Manifest.PathPrefixes != nil { + manifestCopy.PathPrefixes = make([]string, len(src.Manifest.PathPrefixes)) + copy(manifestCopy.PathPrefixes, src.Manifest.PathPrefixes) + } + dst.Manifest = &manifestCopy + } + + // Installation (only scalar fields; pointer dereference suffices) + if src.Installation != nil { + installCopy := *src.Installation + dst.Installation = &installCopy + } + + // ConfigFile (only scalar fields) + if src.ConfigFile != nil { + cfCopy := *src.ConfigFile + dst.ConfigFile = &cfCopy + } + + // Execution (has Args []string and Env map[string]string) + if src.Execution != nil { + execCopy := *src.Execution + if src.Execution.Args != nil { + execCopy.Args = make([]string, len(src.Execution.Args)) + copy(execCopy.Args, src.Execution.Args) + } + if src.Execution.Env != nil { + execCopy.Env = make(map[string]string, len(src.Execution.Env)) + maps.Copy(execCopy.Env, src.Execution.Env) + } + dst.Execution = &execCopy + } + + // MCP (only scalar fields) + if src.MCP != nil { + mcpCopy := *src.MCP + dst.MCP = &mcpCopy + } + + return dst +} + func (c *Compiler) registerNamedEngineDefinitionFromJSON(engineJSON string) error { def, err := parseEngineDefinitionFromJSON(engineJSON) if err != nil || !isEngineDefinitionForm(def) { diff --git a/pkg/workflow/behavior_defined_engine_cache_test.go b/pkg/workflow/behavior_defined_engine_cache_test.go new file mode 100644 index 00000000000..2e4b1a8fe1a --- /dev/null +++ b/pkg/workflow/behavior_defined_engine_cache_test.go @@ -0,0 +1,173 @@ +//go:build !integration + +package workflow + +import ( + "testing" + + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" +) + +// resetEngineDefinitionCacheForTest clears the engine definition cache and +// built-in key set between tests to ensure hermetic behaviour. +func resetEngineDefinitionCacheForTest() { + engineDefinitionCache.Range(func(k, _ any) bool { + engineDefinitionCache.Delete(k) + return true + }) + engineDefinitionBuiltinKeys.Range(func(k, _ any) bool { + engineDefinitionBuiltinKeys.Delete(k) + return true + }) +} + +func TestParseEngineDefinitionFromJSON_EmptyJSON(t *testing.T) { + def, err := parseEngineDefinitionFromJSON("") + assert.Nil(t, def) + assert.NoError(t, err) +} + +func TestParseEngineDefinitionFromJSON_InvalidJSON(t *testing.T) { + _, err := parseEngineDefinitionFromJSON("{bad json") + assert.Error(t, err) +} + +func TestParseEngineDefinitionFromJSON_NonObjectJSON(t *testing.T) { + // A JSON string (not an object) should return nil without error. + def, err := parseEngineDefinitionFromJSON(`"copilot"`) + assert.Nil(t, def) + assert.NoError(t, err) +} + +func TestParseEngineDefinitionFromJSON_ValidObject(t *testing.T) { + const engineJSON = `{"id":"test-engine","display-name":"Test Engine"}` + + def, err := parseEngineDefinitionFromJSON(engineJSON) + require.NoError(t, err) + require.NotNil(t, def) + assert.Equal(t, "test-engine", def.ID) + assert.Equal(t, "Test Engine", def.DisplayName) + // RuntimeID defaults to ID when omitted. + assert.Equal(t, "test-engine", def.RuntimeID) +} + +func TestParseEngineDefinitionFromJSON_CacheHitReturnsCopy(t *testing.T) { + resetEngineDefinitionCacheForTest() + defer resetEngineDefinitionCacheForTest() + + const engineJSON = `{"id":"test-engine","display-name":"Test Engine"}` + + // Pre-seed the built-in key so the engine gets cached. + registerBuiltinEngineDefinitionJSON(engineJSON) + + first, err := parseEngineDefinitionFromJSON(engineJSON) + require.NoError(t, err) + require.NotNil(t, first) + + // Mutate a scalar field on the returned pointer. + first.RuntimeID = "mutated" + + // A second call must return a fresh independent copy. + second, err := parseEngineDefinitionFromJSON(engineJSON) + require.NoError(t, err) + require.NotNil(t, second) + assert.NotEqual(t, "mutated", second.RuntimeID, + "cache hit returned the same object; scalar mutation leaked into cached state") +} + +func TestParseEngineDefinitionFromJSON_CacheHitReturnsDeepCopy(t *testing.T) { + resetEngineDefinitionCacheForTest() + defer resetEngineDefinitionCacheForTest() + + const engineJSON = `{"id":"test-engine","options":{"key":"value"}}` + + registerBuiltinEngineDefinitionJSON(engineJSON) + + first, err := parseEngineDefinitionFromJSON(engineJSON) + require.NoError(t, err) + require.NotNil(t, first) + require.NotNil(t, first.Options) + + // Mutate the Options map on the returned definition. + first.Options["injected"] = "evil" + + second, err := parseEngineDefinitionFromJSON(engineJSON) + require.NoError(t, err) + require.NotNil(t, second) + _, poisoned := second.Options["injected"] + assert.False(t, poisoned, "Options mutation leaked into cached state via shallow map copy") +} + +func TestParseEngineDefinitionFromJSON_NonBuiltinNotCached(t *testing.T) { + resetEngineDefinitionCacheForTest() + defer resetEngineDefinitionCacheForTest() + + const engineJSON = `{"id":"custom-engine"}` + // engineJSON is NOT registered as a built-in key. + + _, err := parseEngineDefinitionFromJSON(engineJSON) + require.NoError(t, err) + + // The cache must remain empty since this is not a known built-in. + found := false + engineDefinitionCache.Range(func(_, _ any) bool { + found = true + return false + }) + assert.False(t, found, "non-builtin engine JSON was incorrectly added to the cache") +} + +func TestDeepCopyEngineDefinition_SlicesAreIndependent(t *testing.T) { + src := EngineDefinition{ + Models: ModelSelection{Supported: []string{"gpt-4", "gpt-3.5"}}, + Auth: []AuthBinding{{Role: "api", Secret: "MY_SECRET"}}, + } + dst := deepCopyEngineDefinition(src) + + // Mutating src slices must not affect dst. + src.Models.Supported[0] = "changed" + src.Auth[0].Secret = "CHANGED" + + assert.Equal(t, "gpt-4", dst.Models.Supported[0]) + assert.Equal(t, "MY_SECRET", dst.Auth[0].Secret) +} + +func TestDeepCopyEngineDefinition_BehaviorsAreIndependent(t *testing.T) { + src := EngineDefinition{ + Behaviors: &EngineBehaviorDefinition{ + SupportedEnvVarKeys: []string{"API_KEY"}, + Execution: &EngineExecutionDefinition{ + Args: []string{"--model", "gpt-4"}, + Env: map[string]string{"FOO": "bar"}, + }, + }, + } + dst := deepCopyEngineDefinition(src) + + // Mutate src's nested reference types. + src.Behaviors.SupportedEnvVarKeys[0] = "CHANGED" + src.Behaviors.Execution.Args[0] = "CHANGED" + src.Behaviors.Execution.Env["FOO"] = "CHANGED" + + require.NotNil(t, dst.Behaviors) + assert.Equal(t, "API_KEY", dst.Behaviors.SupportedEnvVarKeys[0]) + require.NotNil(t, dst.Behaviors.Execution) + assert.Equal(t, "--model", dst.Behaviors.Execution.Args[0]) + assert.Equal(t, "bar", dst.Behaviors.Execution.Env["FOO"]) +} + +func TestDeepCopyAny_NestedMapAndSlice(t *testing.T) { + src := map[string]any{ + "nested": map[string]any{"key": "value"}, + "list": []any{"a", "b"}, + } + dst := deepCopyAny(src).(map[string]any) + + // Mutate src; dst must be unaffected. + src["nested"].(map[string]any)["key"] = "mutated" + src["list"].([]any)[0] = "mutated" + + assert.Equal(t, "value", dst["nested"].(map[string]any)["key"]) + assert.Equal(t, "a", dst["list"].([]any)[0]) +} diff --git a/pkg/workflow/engine_definition_loader.go b/pkg/workflow/engine_definition_loader.go index b230c0945ef..904b844ebde 100644 --- a/pkg/workflow/engine_definition_loader.go +++ b/pkg/workflow/engine_definition_loader.go @@ -20,6 +20,7 @@ package workflow import ( "embed" + "encoding/json" "errors" "fmt" "io/fs" @@ -147,6 +148,21 @@ func loadBuiltinEngineDefinitions() []*EngineDefinition { // file can be resolved and read during import processing. parser.RegisterBuiltinVirtualFile(builtinEnginePath(def.ID), data) + // Register the JSON key that the import processor produces for this built-in + // so that parseEngineDefinitionFromJSON will cache it. The import processor + // marshals the raw "engine:" value (a map[string]any) to JSON; we replicate + // that here so the cache key matches exactly. + var rawFM map[string]any + if jsonErr := yaml.Unmarshal(frontmatterYAML, &rawFM); jsonErr != nil { + engineDefinitionLoaderLog.Printf("Warning: failed to unmarshal frontmatter for cache-key registration of %s: %v", path, jsonErr) + } else if engineVal, ok := rawFM["engine"]; ok { + if jsonKey, jsonErr := json.Marshal(engineVal); jsonErr != nil { + engineDefinitionLoaderLog.Printf("Warning: failed to marshal engine value to JSON for cache-key registration of %s: %v", path, jsonErr) + } else { + registerBuiltinEngineDefinitionJSON(string(jsonKey)) + } + } + engineDefinitionLoaderLog.Printf("Loaded built-in engine definition: id=%s runtime-id=%s", def.ID, def.RuntimeID) definitions = append(definitions, &def) return nil