fixup! Add pi support

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,
 		})
 	}
 

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")
+}

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,

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")

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 {