feat(capabilityclient): forward CTX_AI_MODEL in ai.generate requests

  When CTX_AI_MODEL is set in the container environment (injected by ctx
  from the extension's model: field in ctx.yml), Generate() now includes
  it as a "model" field in the ai.generate POST body. The host capability
  server uses this to create a one-shot provider with the overridden model
  instead of the cached default.

  The field is omitted when CTX_AI_MODEL is unset, preserving existing
  behaviour for callers that do not specify a per-extension model.

Author: miohansen

capabilityclient/generate.go

@@ -4,11 +4,13 @@ import (
 	"context"
 	"encoding/json"
 	"fmt"
+	"os"
 	"strings"
 )
 
 type generateInput struct {
 	Prompt string `json:"prompt"`
+	Model  string `json:"model,omitempty"` // optional per-call model override
 }
 
 type generateOutput struct {
@@ -18,14 +20,22 @@ type generateOutput struct {
 // Generate calls the built-in ai.generate capability with the given prompt and
 // returns the trimmed text output.
 //
+// When CTX_AI_MODEL is set in the container environment (injected by ctx from
+// the extension's model: field in ctx.yml), it is forwarded in the request
+// body so the host capability server uses that model for this call.
+//
 // This is the typed equivalent of:
 //
 //	ctx run ai.generate --input '{"prompt":"..."}'
 func (c *Client) Generate(ctx context.Context, prompt string) (string, error) {
 	if prompt == "" {
 		return "", fmt.Errorf("capability-client: Generate: prompt must not be empty")
 	}
-	raw, err := c.Run(ctx, "ai.generate", generateInput{Prompt: prompt})
+	input := generateInput{Prompt: prompt}
+	if model := os.Getenv("CTX_AI_MODEL"); model != "" {
+		input.Model = model
+	}
+	raw, err := c.Run(ctx, "ai.generate", input)
 	if err != nil {
 		return "", err
 	}

capabilityclient/generate_test.go

@@ -2,7 +2,9 @@ package capabilityclient_test
 
 import (
 	"context"
+	"encoding/json"
 	"errors"
+	"io"
 	"net/http"
 	"net/http/httptest"
 	"testing"
@@ -64,3 +66,70 @@ func TestGenerate_RunError(t *testing.T) {
 		t.Errorf("error type = %T; want *RunError", err)
 	}
 }
+
+func TestGenerate_ForwardsModelFromEnv(t *testing.T) {
+	var capturedBody []byte
+	srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+		capturedBody, _ = io.ReadAll(r.Body)
+		w.Header().Set("Content-Type", "application/json")
+		w.Write([]byte(`{"output":"ok"}`)) //nolint:errcheck
+	}))
+	defer srv.Close()
+	t.Setenv("CTX_CAPABILITY_URL", srv.URL)
+	t.Setenv("CTX_AI_MODEL", "claude-haiku-4-5")
+
+	_, err := capabilityclient.New().Generate(context.Background(), "hello")
+	if err != nil {
+		t.Fatalf("unexpected error: %v", err)
+	}
+
+	// The body is {"input": <inputJSON>}; unwrap to check the model field.
+	var envelope struct {
+		Input json.RawMessage `json:"input"`
+	}
+	if err := json.Unmarshal(capturedBody, &envelope); err != nil {
+		t.Fatalf("parse envelope: %v", err)
+	}
+	var input struct {
+		Model string `json:"model"`
+	}
+	if err := json.Unmarshal(envelope.Input, &input); err != nil {
+		t.Fatalf("parse input: %v", err)
+	}
+	if input.Model != "claude-haiku-4-5" {
+		t.Errorf("model = %q; want claude-haiku-4-5", input.Model)
+	}
+}
+
+func TestGenerate_NoModelWhenEnvUnset(t *testing.T) {
+	var capturedBody []byte
+	srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+		capturedBody, _ = io.ReadAll(r.Body)
+		w.Header().Set("Content-Type", "application/json")
+		w.Write([]byte(`{"output":"ok"}`)) //nolint:errcheck
+	}))
+	defer srv.Close()
+	t.Setenv("CTX_CAPABILITY_URL", srv.URL)
+	t.Setenv("CTX_AI_MODEL", "") // explicitly unset
+
+	_, err := capabilityclient.New().Generate(context.Background(), "hello")
+	if err != nil {
+		t.Fatalf("unexpected error: %v", err)
+	}
+
+	var envelope struct {
+		Input json.RawMessage `json:"input"`
+	}
+	if err := json.Unmarshal(capturedBody, &envelope); err != nil {
+		t.Fatalf("parse envelope: %v", err)
+	}
+	var input struct {
+		Model *string `json:"model"`
+	}
+	if err := json.Unmarshal(envelope.Input, &input); err != nil {
+		t.Fatalf("parse input: %v", err)
+	}
+	if input.Model != nil {
+		t.Errorf("model field present with value %q; want omitted", *input.Model)
+	}
+}