1 files changed, 41 insertions, 0 deletions
diff --git a/libpanto-go/panto_test.go b/libpanto-go/panto_test.go
index 5cfed01..fd9df70 100644
--- a/libpanto-go/panto_test.go
+++ b/libpanto-go/panto_test.go
@@ -128,6 +128,47 @@ func TestMessageBuilderRoundTrip(t *testing.T) {
 	}
 }
 
+// TestMessageBuilderSetIdentity rebuilds a message with a per-message wire
+// identity (and no per-block thinking origin), the resumption path that lets
+// callers avoid copying the identity onto every thinking block.
+func TestMessageBuilderSetIdentity(t *testing.T) {
+	Init()
+	defer Deinit()
+
+	conv, err := NewConversation()
+	if err != nil {
+		t.Fatalf("NewConversation: %v", err)
+	}
+	defer conv.Close()
+
+	ab, err := NewMessageBuilder(RoleAssistant)
+	if err != nil {
+		t.Fatalf("NewMessageBuilder(assistant): %v", err)
+	}
+	// Thinking block with a signature but NO per-block origin.
+	if err := ab.AddThinking("reasoned", "sig", OpenAIChat, "", ""); err != nil {
+		t.Fatalf("AddThinking: %v", err)
+	}
+	if err := ab.AddText("answer"); err != nil {
+		t.Fatalf("AddText: %v", err)
+	}
+	// Provenance is carried at the message level instead.
+	if err := ab.SetIdentity(AnthropicMessages, "https://api.anthropic.com", "claude", ReasoningDefault); err != nil {
+		t.Fatalf("SetIdentity: %v", err)
+	}
+	if err := conv.AddMessage(ab, &Usage{Input: 10, Output: 5}); err != nil {
+		t.Fatalf("AddMessage: %v", err)
+	}
+
+	asst := conv.Message(0).Snapshot()
+	if asst.Role != RoleAssistant || len(asst.Content) != 2 {
+		t.Fatalf("assistant message: role=%v blocks=%d", asst.Role, len(asst.Content))
+	}
+	if b := asst.Content[0]; b.Tag != BlockThinking || b.Text != "reasoned" {
+		t.Fatalf("thinking block: %+v", b)
+	}
+}
+
 func testConfig() Config {
 	return Config{
 		Provider: ProviderConfig{