Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
6 changes: 3 additions & 3 deletions pkg/workflow/compiler_model_pricing.go
Original file line number Diff line number Diff line change
Expand Up @@ -162,9 +162,9 @@ func mergeModelPricingIntoModelCosts(modelCosts map[string]any, provider, model
}
modelEntry := map[string]any{"cost": cost}

// Shallow-clone the top-level map. Do not add +1 here to avoid potential
// integer overflow; Go maps grow automatically as entries are added.
result := make(map[string]any, len(modelCosts))
// Shallow-clone the top-level map without pre-sizing to avoid reintroducing
// allocation-size arithmetic in this security-sensitive path.
result := make(map[string]any)
maps.Copy(result, modelCosts)

// Shallow-clone the providers map.
Expand Down
48 changes: 48 additions & 0 deletions pkg/workflow/compiler_model_pricing_test.go
Original file line number Diff line number Diff line change
Expand Up @@ -4,6 +4,11 @@ package workflow

import (
"context"
"go/ast"
"go/parser"
gotoken "go/token"
"path/filepath"
"runtime"
"testing"

"github.com/stretchr/testify/assert"
Expand Down Expand Up @@ -85,6 +90,49 @@ func TestMergeModelPricingIntoModelCosts_DoesNotMutateInput(t *testing.T) {
assert.Empty(t, openai["models"].(map[string]any))
}

func TestMergeModelPricingIntoModelCosts_DoesNotPreSizeTopLevelClone(t *testing.T) {
_, thisFile, _, ok := runtime.Caller(0)
require.True(t, ok, "runtime.Caller failed")

sourcePath := filepath.Join(filepath.Dir(thisFile), "compiler_model_pricing.go")
file, err := parser.ParseFile(gotoken.NewFileSet(), sourcePath, nil, parser.SkipObjectResolution)
require.NoError(t, err)

var foundResultMake bool
ast.Inspect(file, func(n ast.Node) bool {
fn, ok := n.(*ast.FuncDecl)
if !ok || fn.Name.Name != "mergeModelPricingIntoModelCosts" {
return true
}

for _, stmt := range fn.Body.List {
assign, ok := stmt.(*ast.AssignStmt)
if !ok || len(assign.Lhs) != 1 || len(assign.Rhs) != 1 {
continue
}
lhs, ok := assign.Lhs[0].(*ast.Ident)
if !ok || lhs.Name != "result" {
continue
}
call, ok := assign.Rhs[0].(*ast.CallExpr)
if !ok {
continue
}
fun, ok := call.Fun.(*ast.Ident)
if !ok || fun.Name != "make" {
continue
}

foundResultMake = true
assert.Len(t, call.Args, 1, "result map allocation should not pre-size capacity")
return false
}
return false
Comment on lines +108 to +130
})

assert.True(t, foundResultMake, "expected to find result map allocation in mergeModelPricingIntoModelCosts")
}

// ── resolveEngineProviderForPricing ─────────────────────────────────────────

func TestResolveEngineProviderForPricing(t *testing.T) {
Expand Down
Loading