diff --git a/internal/adapter/common/blocks.go b/internal/adapter/common/blocks.go index e657ce0..b5de3e9 100644 --- a/internal/adapter/common/blocks.go +++ b/internal/adapter/common/blocks.go @@ -31,16 +31,16 @@ func GatewayProviderKey(gatewayURL string) string { } // crushModel builds a 10-field model entry required by the crush schema. -func crushModel(id string) map[string]any { +func crushModel(m canonical.Model) map[string]any { return map[string]any{ - "id": id, - "name": id, + "id": m.ID, + "name": m.ID, "cost_per_1m_in": 0, "cost_per_1m_out": 0, "cost_per_1m_in_cached": 0, "cost_per_1m_out_cached": 0, - "context_window": 200000, - "default_max_tokens": 8192, + "context_window": m.ContextLimit(), + "default_max_tokens": m.OutputLimit(), "can_reason": false, "supports_attachments": true, } @@ -56,10 +56,10 @@ func ProvidersAsCrushMap(p *canonical.Profile) map[string]any { } models := make([]map[string]any, 0) for _, m := range p.Models { - models = append(models, crushModel(m.ID)) + models = append(models, crushModel(m)) } if len(models) == 0 && p.Gateway.DefaultModel != "" { - models = append(models, crushModel(p.Gateway.DefaultModel)) + models = append(models, crushModel(canonical.Model{ID: p.Gateway.DefaultModel})) } out[GatewayProviderKey(p.Gateway.URL)] = map[string]any{ "type": "openai-compat", @@ -101,7 +101,7 @@ func ProvidersAsMap(p *canonical.Profile) map[string]any { } models[m.ID] = map[string]any{ "name": displayName, - "limit": map[string]any{"context": 200000, "output": 8192}, + "limit": map[string]any{"context": m.ContextLimit(), "output": m.OutputLimit()}, } } entry := map[string]any{ @@ -172,7 +172,7 @@ func GooseCustomProviderFile(p *canonical.Profile) ([]byte, string) { for _, m := range p.Models { models = append(models, map[string]any{ "name": m.ID, - "context_limit": 200000, + "context_limit": m.ContextLimit(), }) } if len(models) == 0 && p.Gateway.DefaultModel != "" { @@ -183,7 +183,7 @@ func GooseCustomProviderFile(p *canonical.Profile) ([]byte, string) { } models = append(models, map[string]any{ "name": modelID, - "context_limit": 200000, + "context_limit": canonical.DefaultContextLimit, }) } diff --git a/internal/adapter/common/common_test.go b/internal/adapter/common/common_test.go index 9f10de3..6a9a786 100644 --- a/internal/adapter/common/common_test.go +++ b/internal/adapter/common/common_test.go @@ -1,6 +1,7 @@ package common import ( + "encoding/json" "testing" "github.com/stretchr/testify/assert" @@ -64,3 +65,52 @@ func TestParseMCPFromJSONMissingKey(t *testing.T) { require.NoError(t, err) assert.Nil(t, servers) } + +func TestProvidersAsMapHonorsPerModelLimits(t *testing.T) { + p := &canonical.Profile{ + Gateway: canonical.Gateway{URL: "https://gw.example.com/v1", Token: "t"}, + Models: []canonical.Model{ + {ID: "prov/big", Context: 262144, Output: 65536}, + {ID: "prov/default"}, + }, + } + models := ProvidersAsMap(p)[GatewayProviderKey(p.Gateway.URL)].(map[string]any)["models"].(map[string]any) + + big := models["prov/big"].(map[string]any)["limit"].(map[string]any) + assert.Equal(t, 262144, big["context"], "explicit context override should win") + assert.Equal(t, 65536, big["output"], "explicit output override should win") + + def := models["prov/default"].(map[string]any)["limit"].(map[string]any) + assert.Equal(t, 200000, def["context"], "unset model falls back to default context") + assert.Equal(t, 8192, def["output"], "unset model falls back to default output") +} + +func TestProvidersAsCrushMapHonorsPerModelLimits(t *testing.T) { + p := &canonical.Profile{ + Gateway: canonical.Gateway{URL: "https://gw.example.com/v1", Token: "t"}, + Models: []canonical.Model{ + {ID: "prov/big", Context: 262144, Output: 65536}, + {ID: "prov/default"}, + }, + } + models := ProvidersAsCrushMap(p)[GatewayProviderKey(p.Gateway.URL)].(map[string]any)["models"].([]map[string]any) + + assert.Equal(t, 262144, models[0]["context_window"], "explicit context override should win") + assert.Equal(t, 65536, models[0]["default_max_tokens"], "explicit output override should win") + assert.Equal(t, 200000, models[1]["context_window"], "unset model falls back to default context") + assert.Equal(t, 8192, models[1]["default_max_tokens"], "unset model falls back to default output") +} + +func TestGooseCustomProviderFileHonorsPerModelContext(t *testing.T) { + p := &canonical.Profile{ + Gateway: canonical.Gateway{URL: "https://gw.example.com/v1", Token: "t"}, + Models: []canonical.Model{{ID: "big", Context: 262144}, {ID: "default"}}, + } + body, _ := GooseCustomProviderFile(p) + var entry map[string]any + require.NoError(t, json.Unmarshal(body, &entry)) + models := entry["models"].([]any) + + assert.Equal(t, float64(262144), models[0].(map[string]any)["context_limit"], "explicit context override should win") + assert.Equal(t, float64(200000), models[1].(map[string]any)["context_limit"], "unset model falls back to default context") +} diff --git a/internal/adapters/zed/zed.go b/internal/adapters/zed/zed.go index 1272be4..9f60916 100644 --- a/internal/adapters/zed/zed.go +++ b/internal/adapters/zed/zed.go @@ -122,7 +122,7 @@ func zedLanguageModels(p *canonical.Profile) map[string]any { for _, m := range p.Models { me := map[string]any{ "name": m.ID, - "max_tokens": 200000, + "max_tokens": m.ContextLimit(), "capabilities": map[string]any{ "tools": true, "images": false, diff --git a/internal/adapters/zed/zed_test.go b/internal/adapters/zed/zed_test.go index eff7074..c37108c 100644 --- a/internal/adapters/zed/zed_test.go +++ b/internal/adapters/zed/zed_test.go @@ -63,6 +63,25 @@ func TestZedRenderProducesExpectedTargets(t *testing.T) { assert.NotContains(t, parsed, "context_servers") } +func TestZedRenderHonorsPerModelContext(t *testing.T) { + home := t.TempDir() + b := bundle() + b.Profile.Models = []canonical.Model{{ID: "big", Context: 262144}} + + fs, err := New(WithHome(home)).Render(b) + require.NoError(t, err) + + seen := map[string]adapter.File{} + fs.ForEach(func(f adapter.File) { seen[f.Dest] = f }) + cfgDest := filepath.Join(home, ".config", "zed", "settings.json") + require.Contains(t, seen, cfgDest) + + var parsed map[string]any + require.NoError(t, json.Unmarshal(seen[cfgDest].Content, &parsed)) + models := parsed["language_models"].(map[string]any)["openai_compatible"].(map[string]any)["hs-llmgw-example-com"].(map[string]any)["available_models"].([]any) + assert.Equal(t, float64(262144), models[0].(map[string]any)["max_tokens"], "explicit context override should win") +} + func TestZedRenderMergesExistingKeys(t *testing.T) { home := t.TempDir() base := filepath.Join(home, ".config", "zed") diff --git a/internal/canonical/types.go b/internal/canonical/types.go index 57ca467..8054d3b 100644 --- a/internal/canonical/types.go +++ b/internal/canonical/types.go @@ -35,6 +35,33 @@ type Upstream struct { type Model struct { ID string `yaml:"id"` Alias string `yaml:"alias,omitempty"` + // Context and Output are optional per-model limit overrides (tokens). When + // unset (0), harnesses fall back to DefaultContextLimit / DefaultOutputLimit. + Context int `yaml:"context,omitempty"` + Output int `yaml:"output,omitempty"` +} + +// Default per-model limits, applied by harnesses that surface context/output +// limits when a Model doesn't specify its own. +const ( + DefaultContextLimit = 200000 + DefaultOutputLimit = 8192 +) + +// ContextLimit returns the model's context window, or DefaultContextLimit if unset. +func (m Model) ContextLimit() int { + if m.Context > 0 { + return m.Context + } + return DefaultContextLimit +} + +// OutputLimit returns the model's max output tokens, or DefaultOutputLimit if unset. +func (m Model) OutputLimit() int { + if m.Output > 0 { + return m.Output + } + return DefaultOutputLimit } type MCPRegistry struct {