promptkit/continuation_contract_test.go

package promptkit_test

import (
	"context"
	"encoding/json"
	"os"
	"path/filepath"
	"reflect"
	"sync"
	"testing"

	"go.maximumdirect.net/promptkit"
)

func TestContinuationBranchesAcrossEngines(t *testing.T) {
	ctx := context.Background()
	firstClient := &fakeLLMClient{response: &promptkit.GenerateResponse{Content: "first answer"}}
	firstEngine := newSimpleEngine(t, "first question", promptkit.WithLLMClient(firstClient))
	initial, err := firstEngine.Run(ctx, promptkit.RunRequest{PromptID: "prompt", CaptureExchange: true, APIKey: "synthetic-first-key"})
	if err != nil {
		t.Fatal(err)
	}
	original := initial.Exchange.Details()
	secondClient := &fakeLLMClient{response: &promptkit.GenerateResponse{Content: "follow-up answer"}}
	secondEngine := newNamespacedContractEngine(t, secondClient)
	scoped, err := secondEngine.WithNamespaces(promptkit.Namespaces{Prompt: "alpha", Profile: "alpha"})
	if err != nil {
		t.Fatal(err)
	}
	request := promptkit.RunRequest{PromptID: "shared", ProfileID: "shared", ContinueFrom: initial.Exchange, SessionID: "new-session", Validation: &promptkit.OutputContract{Format: promptkit.FormatJSON, ValidationMode: promptkit.ValidationJSON}, AppendedMessages: []promptkit.RenderedMessage{{Role: promptkit.RoleUser, Content: "branch one"}}}
	branch, err := scoped.Run(ctx, request)
	if err != nil {
		t.Fatal(err)
	}
	request.AppendedMessages[0].Content = "branch two"
	sibling, err := scoped.Run(ctx, request)
	if err != nil {
		t.Fatal(err)
	}
	next, err := scoped.Run(ctx, promptkit.RunRequest{PromptID: "shared", ProfileID: "shared", ContinueFrom: branch.Exchange})
	if err != nil {
		t.Fatal(err)
	}
	if branch.ParentRunID != initial.RunID || sibling.ParentRunID != initial.RunID || next.ParentRunID != branch.RunID || branch.PromptID != "alpha/shared" || branch.SelectedProfileID != "alpha/shared" || branch.SessionID != "new-session" {
		t.Fatalf("ancestry or selections incorrect: %#v %#v %#v", branch, sibling, next)
	}
	if !reflect.DeepEqual(initial.Exchange.Details(), original) {
		t.Fatal("branching changed source")
	}
	if branch.Validation.Status != promptkit.ValidationFailed || branch.ModelName != "alpha-model" {
		t.Fatalf("follow-up settings: %#v", branch)
	}
	branchMessages := branch.Exchange.Details().Messages
	siblingMessages := sibling.Exchange.Details().Messages
	if len(branchMessages) != 5 || len(siblingMessages) != 5 || branchMessages[3].Content != "branch one" || siblingMessages[3].Content != "branch two" || len(next.Exchange.Details().Messages) != 7 {
		t.Fatalf("branch histories = %#v / %#v", branchMessages, siblingMessages)
	}
	encoded, err := json.Marshal(branch.Exchange.Details())
	if err != nil || !json.Valid(encoded) {
		t.Fatalf("details JSON = %s, %v", encoded, err)
	}
	branchMessages[0].Content = "changed details"
	branch.RawOutput = "changed result"
	if initial.Exchange.Details().Messages[0].Content != "first question" || branch.Exchange.Details().Messages[4].Content != "follow-up answer" {
		t.Fatal("caller mutation changed captured history")
	}
	for _, request := range secondClient.requests {
		if request.APIKey == "synthetic-first-key" || request.Prompt.Messages[0].Content != "first question" || request.Prompt.Messages[1].Content != "first answer" {
			t.Fatalf("unexpected follow-up request: %#v", request)
		}
	}
}

func TestContinuationDoesNotReloadCapturedInputs(t *testing.T) {
	ctx := context.Background()
	path := filepath.Join(t.TempDir(), "input.txt")
	if err := os.WriteFile(path, []byte("original input"), 0600); err != nil {
		t.Fatal(err)
	}
	client := &fakeLLMClient{response: &promptkit.GenerateResponse{Content: ""}}
	engine := newSimpleEngine(t, `{{input "source"}}`, promptkit.WithLLMClient(client))
	initial, err := engine.Run(ctx, promptkit.RunRequest{PromptID: "prompt", Inputs: map[string]promptkit.ArtifactRef{"source": promptkit.File(path)}, CaptureExchange: true})
	if err != nil {
		t.Fatal(err)
	}
	if err := os.WriteFile(path, []byte("new input"), 0600); err != nil {
		t.Fatal(err)
	}
	request := promptkit.RunRequest{PromptID: "prompt", Inputs: map[string]promptkit.ArtifactRef{"source": promptkit.File(path)}, ContinueFrom: initial.Exchange}
	prepared, err := engine.PrepareExecution(ctx, request)
	if err != nil {
		t.Fatal(err)
	}
	if err := os.Remove(path); err != nil {
		t.Fatal(err)
	}
	result, err := engine.RunPrepared(ctx, prepared)
	if err != nil {
		t.Fatal(err)
	}
	want := []promptkit.RenderedMessage{{Role: promptkit.RoleUser, Content: "original input"}, {Role: promptkit.RoleAssistant, Content: ""}, {Role: promptkit.RoleUser, Content: "new input"}, {Role: promptkit.RoleAssistant, Content: ""}}
	if !reflect.DeepEqual(result.Exchange.Details().Messages, want) {
		t.Fatalf("history = %#v", result.Exchange.Details().Messages)
	}
	if result, err := engine.Run(ctx, request); err == nil || result != nil {
		t.Fatalf("missing new input = %#v, %v", result, err)
	}
}

func TestContinuationSourceSupportsConcurrentRuns(t *testing.T) {
	client := &preparedRecordingClient{response: &promptkit.GenerateResponse{Content: "answer"}}
	engine := newSimpleEngine(t, "question", promptkit.WithLLMClient(client))
	initial, err := engine.Run(context.Background(), promptkit.RunRequest{PromptID: "prompt", CaptureExchange: true})
	if err != nil {
		t.Fatal(err)
	}
	const count = 8
	var wg sync.WaitGroup
	errors := make(chan error, count)
	for i := 0; i < count; i++ {
		wg.Add(1)
		go func() {
			defer wg.Done()
			result, err := engine.Run(context.Background(), promptkit.RunRequest{PromptID: "prompt", ContinueFrom: initial.Exchange})
			if err == nil && (result.Exchange == nil || result.ParentRunID != initial.RunID) {
				t.Error("missing continuation ancestry")
			}
			errors <- err
		}()
	}
	wg.Wait()
	close(errors)
	for err := range errors {
		if err != nil {
			t.Fatal(err)
		}
	}
	if len(initial.Exchange.Details().Messages) != 2 {
		t.Fatal("source changed")
	}
}