Vectrum
Sign In
124 lines · 4.1 KB · Go
package llm

import (
	"encoding/json"
	"strings"
	"testing"
)

func TestContentBlockMarshalsOnlyItsOwnFields(t *testing.T) {
	tests := []struct {
		name  string
		block ContentBlock
		want  string
	}{
		{
			name:  "text carries no stray thinking key",
			block: TextBlock("hello"),
			want:  `{"text":"hello","type":"text"}`,
		},
		{
			name:  "tool_use keeps input as raw JSON",
			block: ContentBlock{Type: BlockToolUse, ID: "tu_1", Name: "read_file", Input: json.RawMessage(`{"path":"main.go"}`)},
			want:  `{"id":"tu_1","input":{"path":"main.go"},"name":"read_file","type":"tool_use"}`,
		},
		{
			name:  "tool_use with no input sends an empty object",
			block: ContentBlock{Type: BlockToolUse, ID: "tu_2", Name: "pwd"},
			want:  `{"id":"tu_2","input":{},"name":"pwd","type":"tool_use"}`,
		},
		{
			name:  "successful tool_result omits is_error",
			block: ToolResultBlock("tu_1", "ok", false),
			want:  `{"content":"ok","tool_use_id":"tu_1","type":"tool_result"}`,
		},
		{
			name:  "failed tool_result sets is_error",
			block: ToolResultBlock("tu_1", "no such file", true),
			want:  `{"content":"no such file","is_error":true,"tool_use_id":"tu_1","type":"tool_result"}`,
		},
	}
	for _, tc := range tests {
		t.Run(tc.name, func(t *testing.T) {
			got, err := json.Marshal(tc.block)
			if err != nil {
				t.Fatalf("marshal: %v", err)
			}
			if string(got) != tc.want {
				t.Errorf("got  %s\nwant %s", got, tc.want)
			}
		})
	}
}

// With display omitted (the default on Opus 5) a thinking block has empty text
// but a signature that must go back unchanged, so the field cannot be omitempty.
func TestThinkingBlockRoundTripsEmptyTextAndSignature(t *testing.T) {
	const wire = `{"type":"thinking","thinking":"","signature":"sig-abc123"}`
	var b ContentBlock
	if err := json.Unmarshal([]byte(wire), &b); err != nil {
		t.Fatalf("unmarshal: %v", err)
	}
	out, err := json.Marshal(b)
	if err != nil {
		t.Fatalf("marshal: %v", err)
	}
	if !strings.Contains(string(out), `"thinking":""`) {
		t.Errorf("empty thinking text must survive the round trip, got %s", out)
	}
	if !strings.Contains(string(out), `"signature":"sig-abc123"`) {
		t.Errorf("signature must survive the round trip, got %s", out)
	}
}

func TestUnknownBlockTypeRoundTripsVerbatim(t *testing.T) {
	const wire = `{"type":"server_tool_use_future","id":"x","weird":{"nested":[1,2]}}`
	var b ContentBlock
	if err := json.Unmarshal([]byte(wire), &b); err != nil {
		t.Fatalf("unmarshal: %v", err)
	}
	out, err := json.Marshal(b)
	if err != nil {
		t.Fatalf("marshal: %v", err)
	}
	if string(out) != wire {
		t.Errorf("unmodelled block should be echoed verbatim\ngot  %s\nwant %s", out, wire)
	}
}

func TestEmptyToolResultGetsPlaceholder(t *testing.T) {
	if got := ToolResultBlock("tu_1", "", false).Content; got != "(no output)" {
		t.Errorf("content = %q, want a placeholder: the API rejects empty tool_result content", got)
	}
}

func TestResponseAccessors(t *testing.T) {
	const wire = `{
	  "id":"msg_1","type":"message","role":"assistant","model":"claude-opus-5",
	  "content":[
	    {"type":"thinking","thinking":"","signature":"s"},
	    {"type":"text","text":"Let me look. "},
	    {"type":"tool_use","id":"tu_1","name":"read_file","input":{"path":"a.go"}},
	    {"type":"text","text":"One moment."}
	  ],
	  "stop_reason":"tool_use","stop_details":null,
	  "usage":{"input_tokens":10,"output_tokens":20,"cache_read_input_tokens":5,"cache_creation_input_tokens":1}
	}`
	var r Response
	if err := json.Unmarshal([]byte(wire), &r); err != nil {
		t.Fatalf("unmarshal: %v", err)
	}
	if got := r.Text(); got != "Let me look. One moment." {
		t.Errorf("Text() = %q, want the text blocks concatenated", got)
	}
	uses := r.ToolUses()
	if len(uses) != 1 || uses[0].ID != "tu_1" || uses[0].Name != "read_file" {
		t.Errorf("ToolUses() = %+v, want the single read_file call", uses)
	}
	if r.Usage.CacheReadInputTokens != 5 {
		t.Errorf("cache_read_input_tokens = %d, want 5", r.Usage.CacheReadInputTokens)
	}
	if msg := r.Message(); msg.Role != RoleAssistant || len(msg.Content) != 4 {
		t.Errorf("Message() = %+v, want all 4 blocks under the assistant role", msg)
	}
}