1 files changed, 21 insertions, 39 deletions
diff --git a/libpanto-go/panto_test.go b/libpanto-go/panto_test.go
index fd9df70..af8cafc 100644
--- a/libpanto-go/panto_test.go
+++ b/libpanto-go/panto_test.go
@@ -59,7 +59,7 @@ func TestCustomStoreMintsSession(t *testing.T) {
 }
 
 // Gap #3: rebuild a conversation containing tool_use / tool_result / thinking
-// via MessageBuilder, then read the blocks back losslessly.
+// via the value-block API, then read the blocks back losslessly.
 func TestMessageBuilderRoundTrip(t *testing.T) {
 	Init()
 	defer Deinit()
@@ -70,31 +70,17 @@ func TestMessageBuilderRoundTrip(t *testing.T) {
 	}
 	defer conv.Close()
 
-	ub, err := NewMessageBuilder(RoleUser)
-	if err != nil {
-		t.Fatalf("NewMessageBuilder(user): %v", err)
-	}
-	if err := ub.AddToolResult("call-1", "result text", false); err != nil {
-		t.Fatalf("AddToolResult: %v", err)
-	}
-	if err := conv.AddMessage(ub, nil); err != nil {
+	if err := conv.AddMessage(UserMessage(
+		ToolResultBlock("call-1", "result text", false),
+	)); err != nil {
 		t.Fatalf("AddMessage(user): %v", err)
 	}
 
-	ab, err := NewMessageBuilder(RoleAssistant)
-	if err != nil {
-		t.Fatalf("NewMessageBuilder(assistant): %v", err)
-	}
-	if err := ab.AddText("let me search"); err != nil {
-		t.Fatalf("AddText: %v", err)
-	}
-	if err := ab.AddThinking("hmm", "sig", AnthropicMessages, "https://api", "claude"); err != nil {
-		t.Fatalf("AddThinking: %v", err)
-	}
-	if err := ab.AddToolUse("call-1", "search", `{"q":1}`); err != nil {
-		t.Fatalf("AddToolUse: %v", err)
-	}
-	if err := conv.AddMessage(ab, &Usage{Input: 10, Output: 5}); err != nil {
+	if err := conv.AddMessage(AssistantMessage(&Usage{Input: 10, Output: 5},
+		TextBlock("let me search"),
+		ThinkingBlock("hmm", "sig", AnthropicMessages, "https://api", "claude"),
+		ToolUseBlock("call-1", "search", `{"q":1}`),
+	)); err != nil {
 		t.Fatalf("AddMessage(assistant): %v", err)
 	}
 
@@ -141,22 +127,18 @@ func TestMessageBuilderSetIdentity(t *testing.T) {
 	}
 	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 {
+	// Thinking block with a signature but NO per-block origin; provenance is
+	// carried at the message level instead, via WithIdentity.
+	msg := AssistantMessage(&Usage{Input: 10, Output: 5},
+		ThinkingBlock("reasoned", "sig", OpenAIChat, "", ""),
+		TextBlock("answer"),
+	).WithIdentity(MessageIdentity{
+		APIStyle:  AnthropicMessages,
+		BaseURL:   "https://api.anthropic.com",
+		Model:     "claude",
+		Reasoning: ReasoningDefault,
+	})
+	if err := conv.AddMessage(msg); err != nil {
 		t.Fatalf("AddMessage: %v", err)
 	}