Adapt corrections through PromptKit

internal/framework/llm/promptkit_client_test.go

@@ -1,6 +1,7 @@
 package llm
 
 import (
+	"bytes"
 	"context"
 	"encoding/json"
 	"errors"
@@ -116,6 +117,74 @@ func TestPromptKitClientMapsPromptRequestAndUnmarshalsOutput(t *testing.T) {
 	}
 }
 
+func TestPromptKitClientAppendsSemanticCorrectionAfterRenderedPrompt(t *testing.T) {
+	fake := &fakePromptKitLLM{content: `{"ok":true}`}
+	client := newTestPromptKitClient(t, fake)
+	request := contracts.StructuredCompletionRequest{
+		PromptID:  "adapter.direct-session",
+		ProfileID: "explicit-profile",
+		SessionID: "correction-session",
+		Inputs: contracts.LLMInputSet{
+			"transcript": contracts.NewLLMInputMaterial("transcript", "application/json", []byte(`{"source":true}`), "", ""),
+		},
+		Vars: map[string]any{"custom": "value"},
+	}
+
+	var ordinary map[string]any
+	if _, err := client.CompleteStructured(context.Background(), request, &ordinary); err != nil {
+		t.Fatalf("ordinary CompleteStructured() error = %v", err)
+	}
+	ordinaryMessages := append([]promptkit.RenderedMessage(nil), fake.lastRequest().Prompt.Messages...)
+	if len(ordinaryMessages) != 1 {
+		t.Fatalf("ordinary rendered messages = %#v, want only the declared prompt message", ordinaryMessages)
+	}
+
+	correction, err := contracts.NewSemanticCorrection([]byte(`{"previous":"response"}`), "Return the corrected JSON object.")
+	if err != nil {
+		t.Fatalf("NewSemanticCorrection() error = %v", err)
+	}
+	request.Correction = correction
+	var corrected map[string]any
+	if _, err := client.CompleteStructured(context.Background(), request, &corrected); err != nil {
+		t.Fatalf("corrected CompleteStructured() error = %v", err)
+	}
+	correctedMessages := fake.lastRequest().Prompt.Messages
+	if len(correctedMessages) != len(ordinaryMessages)+2 {
+		t.Fatalf("corrected message count = %d, want %d", len(correctedMessages), len(ordinaryMessages)+2)
+	}
+	if !reflect.DeepEqual(correctedMessages[:len(ordinaryMessages)], ordinaryMessages) {
+		t.Fatalf("ordinary rendered prefix changed: got %#v, want %#v", correctedMessages[:len(ordinaryMessages)], ordinaryMessages)
+	}
+	if got, want := correctedMessages[len(ordinaryMessages)], (promptkit.RenderedMessage{Role: promptkit.RoleAssistant, Content: `{"previous":"response"}`}); got != want {
+		t.Fatalf("assistant correction message = %#v, want %#v", got, want)
+	}
+	if got, want := correctedMessages[len(ordinaryMessages)+1], (promptkit.RenderedMessage{Role: promptkit.RoleUser, Content: "Return the corrected JSON object."}); got != want {
+		t.Fatalf("user correction message = %#v, want %#v", got, want)
+	}
+}
+
+func TestPromptKitClientRejectsInvalidCorrectionsBeforePromptPreparation(t *testing.T) {
+	const sensitiveResponse = "assistant-response-must-not-appear-in-errors"
+	for _, correction := range []*contracts.SemanticCorrection{
+		{AssistantResponse: []byte(sensitiveResponse), UserGuidance: " \t"},
+		{AssistantResponse: bytes.Repeat([]byte(sensitiveResponse), contracts.MaxAssistantResponseBytes/len(sensitiveResponse)+1), UserGuidance: "Use a smaller response."},
+	} {
+		fake := &fakePromptKitLLM{content: `{"ok":true}`}
+		client := newTestPromptKitClient(t, fake)
+		var out map[string]any
+		_, err := client.CompleteStructured(context.Background(), contracts.StructuredCompletionRequest{Correction: correction}, &out)
+		if err == nil || !strings.Contains(err.Error(), "correction") {
+			t.Fatalf("CompleteStructured() error = %v, want correction validation failure", err)
+		}
+		if strings.Contains(err.Error(), sensitiveResponse) {
+			t.Fatalf("correction validation error leaked response content: %v", err)
+		}
+		if got := atomic.LoadInt32(&fake.calls); got != 0 {
+			t.Fatalf("provider calls = %d, want no provider call after invalid correction", got)
+		}
+	}
+}
+
 func TestPromptKitClientUsesOnePreparedSnapshotForDebugAndGeneration(t *testing.T) {
 	const initialPrompt = `id: snapshot.test
 version: "v1"
@@ -960,6 +1029,10 @@ func TestPromptKitClientRepairsStructuredOutputAndReportsCumulativeUsage(t *test
 		{Content: `{"ok":true}`, Usage: promptkit.TokenUsage{PromptTokens: 7, CompletionTokens: 11, TotalTokens: 18}},
 	}}
 	client := newTestPromptKitClient(t, fake)
+	correction, err := contracts.NewSemanticCorrection([]byte(`{"bad":true}`), "Return the required ok field.")
+	if err != nil {
+		t.Fatalf("NewSemanticCorrection() error = %v", err)
+	}
 
 	var out struct {
 		OK bool `json:"ok"`
@@ -968,6 +1041,7 @@ func TestPromptKitClientRepairsStructuredOutputAndReportsCumulativeUsage(t *test
 		PromptID:                       "adapter.test",
 		StructuredOutputRepairAttempts: &attempts,
 		SessionID:                      "repair-test",
+		Correction:                     correction,
 		Inputs: contracts.LLMInputSet{
 			"transcript": contracts.NewLLMInputMaterial("transcript", "application/json", []byte(`{"source":true}`), "", ""),
 		},
@@ -978,6 +1052,17 @@ func TestPromptKitClientRepairsStructuredOutputAndReportsCumulativeUsage(t *test
 	if got := atomic.LoadInt32(&fake.calls); got != 2 {
 		t.Fatalf("provider calls = %d, want initial generation and one repair", got)
 	}
+	requests := fake.requestsSnapshot()
+	if len(requests) != 2 || len(requests[0].Prompt.Messages) < 3 {
+		t.Fatalf("repair requests = %#v, want correction messages on the initial prepared request", requests)
+	}
+	messages := requests[0].Prompt.Messages
+	if got, want := messages[len(messages)-2], (promptkit.RenderedMessage{Role: promptkit.RoleAssistant, Content: `{"bad":true}`}); got != want {
+		t.Fatalf("repair assistant correction = %#v, want %#v", got, want)
+	}
+	if got, want := messages[len(messages)-1], (promptkit.RenderedMessage{Role: promptkit.RoleUser, Content: "Return the required ok field."}); got != want {
+		t.Fatalf("repair user correction = %#v, want %#v", got, want)
+	}
 	if response.RepairAttempts != 1 || response.PromptTokens != 10 || response.CompletionTokens != 16 || response.TotalTokens != 26 {
 		t.Fatalf("response repair and usage = %#v, want one repair and PromptKit cumulative usage", response)
 	}