Add engine-scoped backend registration

2026-07-29 17:15:51 +00:00
parent 810f80e7c9
commit ae210b3c26
6 changed files with 318 additions and 9 deletions
--- a/public_contract_test.go
+++ b/public_contract_test.go
@@ -83,6 +83,232 @@ func TestUnknownProfileBackendHasProfileLoadIdentity(t *testing.T) {
 	}
 }

+func TestCustomBackendFlowsThroughProfilesOverridesAndInjectedClient(t *testing.T) {
+	t.Setenv("CUSTOM_LLM_KEY", "test-key")
+	client := &fakeLLMClient{response: &promptkit.GenerateResponse{Content: "ok"}}
+	engine, err := promptkit.NewEngine(promptkit.Config{},
+		promptkit.WithPromptFS(contractPromptFS("prompt", "backend-profile", "message"), "."),
+		promptkit.WithBackend(promptkit.Backend{
+			ID:        "  custom  ",
+			Endpoint:  "  http://backend.example/v1  ",
+			APIKeyEnv: "  CUSTOM_LLM_KEY  ",
+			ExtraParams: map[string]any{
+				"provider": "custom",
+			},
+		}),
+		promptkit.WithProfiles(
+			promptkit.Profile{ID: "backend-profile", BackendID: "custom", Model: "backend-model"},
+			promptkit.Profile{ID: "profile-endpoint", BackendID: "custom", Endpoint: "http://profile.example/v1", Model: "profile-model"},
+		),
+		promptkit.WithLLMClient(client),
+	)
+	if err != nil {
+		t.Fatalf("construct engine: %v", err)
+	}
+
+	result, err := engine.Run(context.Background(), promptkit.RunRequest{PromptID: "prompt"})
+	if err != nil {
+		t.Fatalf("run with custom backend: %v", err)
+	}
+	if len(client.requests) != 1 {
+		t.Fatalf("expected one injected-client request, got %d", len(client.requests))
+	}
+	target := client.requests[0].Target
+	if target.BackendID != "custom" ||
+		target.Endpoint != "http://backend.example/v1" ||
+		target.APIKeyEnv != "CUSTOM_LLM_KEY" ||
+		target.Model != "backend-model" ||
+		target.ExtraParams["provider"] != "custom" ||
+		result.SelectedBackendID != "custom" {
+		t.Fatalf("unexpected custom backend settings: target=%+v result_backend=%q", target, result.SelectedBackendID)
+	}
+
+	prepared, err := engine.Prepare(context.Background(), promptkit.RunRequest{
+		PromptID: "prompt", ProfileID: "profile-endpoint",
+	})
+	if err != nil {
+		t.Fatalf("prepare profile endpoint override: %v", err)
+	}
+	if prepared.SelectedBackendID != "custom" || prepared.EffectiveModelParams.Endpoint != "http://profile.example/v1" {
+		t.Fatalf("profile endpoint override changed backend identity: %+v", prepared)
+	}
+
+	prepared, err = engine.Prepare(context.Background(), promptkit.RunRequest{
+		PromptID: "prompt",
+		Execution: &promptkit.ExecutionTargetOverride{
+			Endpoint: "http://request.example/v1",
+		},
+	})
+	if err != nil {
+		t.Fatalf("prepare request endpoint override: %v", err)
+	}
+	if prepared.SelectedBackendID != "custom" || prepared.EffectiveModelParams.Endpoint != "http://request.example/v1" {
+		t.Fatalf("request endpoint override changed backend identity: %+v", prepared)
+	}
+}
+
+func TestCustomBackendSupportsFileProfileAndBothSelectionPaths(t *testing.T) {
+	engine, err := promptkit.NewEngine(promptkit.Config{},
+		promptkit.WithPromptFS(contractPromptFS("prompt", "file-profile", "message"), "."),
+		promptkit.WithProfileFS(fstest.MapFS{
+			"profile.yaml": &fstest.MapFile{Data: []byte(`id: file-profile
+backend: file-backend
+model: file-model
+`)},
+		}, "."),
+		promptkit.WithBackend(promptkit.Backend{
+			ID:       "file-backend",
+			Endpoint: "http://file-backend.example/v1",
+		}),
+	)
+	if err != nil {
+		t.Fatalf("construct engine: %v", err)
+	}
+
+	for _, request := range []promptkit.RunRequest{
+		{PromptID: "prompt"},
+		{PromptID: "prompt", ProfileID: "file-profile"},
+	} {
+		prepared, err := engine.Prepare(context.Background(), request)
+		if err != nil {
+			t.Fatalf("prepare file profile: %v", err)
+		}
+		if prepared.SelectedBackendID != "file-backend" ||
+			prepared.EffectiveModelParams.Endpoint != "http://file-backend.example/v1" {
+			t.Fatalf("unexpected file-profile backend resolution: %+v", prepared)
+		}
+	}
+}
+
+func TestBackendOptionsAccumulateAndRegistrationsAreEngineLocal(t *testing.T) {
+	engine, err := promptkit.NewEngine(promptkit.Config{},
+		promptkit.WithPromptFS(contractPromptFS("prompt", "first-profile", "message"), "."),
+		promptkit.WithBackend(promptkit.Backend{ID: "first", Endpoint: "http://first.example/v1"}),
+		promptkit.WithBackend(promptkit.Backend{ID: "second", Endpoint: "http://second.example/v1"}),
+		promptkit.WithProfiles(
+			promptkit.Profile{ID: "first-profile", BackendID: "first", Model: "model"},
+			promptkit.Profile{ID: "second-profile", BackendID: "second", Model: "model"},
+		),
+	)
+	if err != nil {
+		t.Fatalf("construct engine with accumulated registrations: %v", err)
+	}
+	for profileID, wantEndpoint := range map[string]string{
+		"first-profile":  "http://first.example/v1",
+		"second-profile": "http://second.example/v1",
+	} {
+		prepared, err := engine.Prepare(context.Background(), promptkit.RunRequest{
+			PromptID: "prompt", ProfileID: profileID,
+		})
+		if err != nil {
+			t.Fatalf("prepare %s: %v", profileID, err)
+		}
+		if prepared.EffectiveModelParams.Endpoint != wantEndpoint {
+			t.Fatalf("profile %s endpoint=%q, want %q", profileID, prepared.EffectiveModelParams.Endpoint, wantEndpoint)
+		}
+	}
+
+	newEngine := func(endpoint string) *promptkit.Engine {
+		t.Helper()
+		value, err := promptkit.NewEngine(promptkit.Config{},
+			promptkit.WithPromptFS(contractPromptFS("prompt", "profile", "message"), "."),
+			promptkit.WithBackend(promptkit.Backend{ID: "same-id", Endpoint: endpoint}),
+			promptkit.WithProfiles(promptkit.Profile{ID: "profile", BackendID: "same-id", Model: "model"}),
+		)
+		if err != nil {
+			t.Fatalf("construct isolated engine: %v", err)
+		}
+		return value
+	}
+	firstEngine := newEngine("http://one.example/v1")
+	secondEngine := newEngine("http://two.example/v1")
+	for engine, wantEndpoint := range map[*promptkit.Engine]string{
+		firstEngine:  "http://one.example/v1",
+		secondEngine: "http://two.example/v1",
+	} {
+		prepared, err := engine.Prepare(context.Background(), promptkit.RunRequest{PromptID: "prompt"})
+		if err != nil {
+			t.Fatalf("prepare isolated engine: %v", err)
+		}
+		if prepared.EffectiveModelParams.Endpoint != wantEndpoint {
+			t.Fatalf("isolated engine endpoint=%q, want %q", prepared.EffectiveModelParams.Endpoint, wantEndpoint)
+		}
+	}
+}
+
+func TestBackendRegistrationRejectsInvalidAndDuplicateDefinitions(t *testing.T) {
+	cycle := map[string]any{}
+	cycle["self"] = cycle
+	tests := []struct {
+		name     string
+		backends []promptkit.Backend
+	}{
+		{name: "blank id", backends: []promptkit.Backend{{Endpoint: "http://example.test/v1"}}},
+		{name: "invalid endpoint", backends: []promptkit.Backend{{ID: "custom", Endpoint: "ftp://example.test/v1"}}},
+		{name: "invalid environment", backends: []promptkit.Backend{{ID: "custom", Endpoint: "http://example.test/v1", APIKeyEnv: "BAD-NAME"}}},
+		{name: "reserved extra parameter", backends: []promptkit.Backend{{ID: "custom", Endpoint: "http://example.test/v1", ExtraParams: map[string]any{"model": "override"}}}},
+		{name: "cyclic extra parameter", backends: []promptkit.Backend{{ID: "custom", Endpoint: "http://example.test/v1", ExtraParams: cycle}}},
+		{name: "duplicate consumer id", backends: []promptkit.Backend{
+			{ID: " custom ", Endpoint: "http://one.example/v1"},
+			{ID: "custom", Endpoint: "http://two.example/v1"},
+		}},
+		{name: "reserved built-in id", backends: []promptkit.Backend{{
+			ID: promptkit.BackendOpenRouter, Endpoint: "http://replacement.example/v1",
+		}}},
+	}
+
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			options := []promptkit.Option{
+				promptkit.WithPromptFS(contractPromptFS("prompt", "profile", "message"), "."),
+			}
+			for _, backend := range tt.backends {
+				options = append(options, promptkit.WithBackend(backend))
+			}
+			_, err := promptkit.NewEngine(promptkit.Config{}, options...)
+			if !errors.Is(err, promptkit.ErrInvalidConfig) {
+				t.Fatalf("expected ErrInvalidConfig, got %v", err)
+			}
+		})
+	}
+}
+
+func TestBackendExtraParamsAreDeeplyCopiedAtConstructionAndLookup(t *testing.T) {
+	nested := map[string]any{"value": "original"}
+	extraParams := map[string]any{"nested": nested}
+	engine, err := promptkit.NewEngine(promptkit.Config{},
+		promptkit.WithPromptFS(contractPromptFS("prompt", "profile", "message"), "."),
+		promptkit.WithBackend(promptkit.Backend{
+			ID: "custom", Endpoint: "http://example.test/v1", ExtraParams: extraParams,
+		}),
+		promptkit.WithProfiles(promptkit.Profile{ID: "profile", BackendID: "custom", Model: "model"}),
+	)
+	if err != nil {
+		t.Fatalf("construct engine: %v", err)
+	}
+	nested["value"] = "mutated input"
+	extraParams["later"] = true
+
+	prepared, err := engine.Prepare(context.Background(), promptkit.RunRequest{PromptID: "prompt"})
+	if err != nil {
+		t.Fatalf("first prepare: %v", err)
+	}
+	gotNested := prepared.EffectiveModelParams.ExtraParams["nested"].(map[string]any)
+	if gotNested["value"] != "original" || prepared.EffectiveModelParams.ExtraParams["later"] != nil {
+		t.Fatalf("backend retained caller mutations: %#v", prepared.EffectiveModelParams.ExtraParams)
+	}
+	gotNested["value"] = "mutated lookup"
+
+	prepared, err = engine.Prepare(context.Background(), promptkit.RunRequest{PromptID: "prompt"})
+	if err != nil {
+		t.Fatalf("second prepare: %v", err)
+	}
+	gotNested = prepared.EffectiveModelParams.ExtraParams["nested"].(map[string]any)
+	if gotNested["value"] != "original" {
+		t.Fatalf("backend retained lookup mutation: %#v", prepared.EffectiveModelParams.ExtraParams)
+	}
+}
+
 func TestPreparedRunJSONTimingRoundTrips(t *testing.T) {
 	start := time.Date(2026, time.July, 29, 12, 0, 0, 0, time.UTC)
 	prepared := promptkit.PreparedRun{