Vectrum
Sign In
289 lines · 9.7 KB · Go
package serve

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

	"codey/internal/acp"
)

// The handshake T3 Code actually sends, byte for byte.
//
// These are not paraphrases of the protocol. They are the literal payloads
// built in T3 Code's ACP session runtime, and the assertions below are its
// schema's required fields. When a nightly of theirs breaks us, this is the
// test that says where — which is worth more than a test written against what
// the specification says ought to happen.
//
// Source: apps/server/src/provider/acp/AcpSessionRuntime.ts, at t3code
// 9accc567. initialize at :696-705 with capabilities from :566-578,
// authenticate at :710-718, session/new at :836-849, session/prompt at
// :1005-1014.
const (
	t3Initialize   = `{"protocolVersion":1,"clientCapabilities":{"fs":{"readTextFile":false,"writeTextFile":false},"terminal":false},"clientInfo":{"name":"t3-code","version":"0.0.0"}}`
	t3Authenticate = `{"methodId":"cached_token"}`
)

// raw sends a method with params already encoded, so nothing between the test
// and the wire can quietly fix a payload up.
func (c *testClient) raw(method, params string) json.RawMessage {
	c.t.Helper()
	c.nextID++
	id := json.RawMessage(itoa(c.nextID))
	c.write(acpMessage{JSONRPC: "2.0", ID: &id, Method: method, Params: json.RawMessage(params)})

	m := c.pump(string(id))
	if m.Error != nil {
		c.t.Fatalf("%s: %v", method, m.Error)
	}
	return m.Result
}

// acpMessage mirrors the wire frame so this file can write one by hand.
type acpMessage struct {
	JSONRPC string           `json:"jsonrpc"`
	ID      *json.RawMessage `json:"id,omitempty"`
	Method  string           `json:"method,omitempty"`
	Params  json.RawMessage  `json:"params,omitempty"`
}

func field(t *testing.T, raw json.RawMessage, name string) any {
	t.Helper()
	var m map[string]any
	if err := json.Unmarshal(raw, &m); err != nil {
		t.Fatalf("result is not an object: %s", raw)
	}
	v, ok := m[name]
	if !ok {
		t.Fatalf("result has no %q: %s", name, raw)
	}
	return v
}

func TestTheHandshakeT3CodeActuallySends(t *testing.T) {
	f := newFakeAPI(t, msg("end_turn", textBlock("hello from codey")))
	c := newTestServer(t, f, Options{})

	// initialize — the client requires protocolVersion, a number.
	init := c.raw("initialize", t3Initialize)
	if v, ok := field(t, init, "protocolVersion").(float64); !ok || int(v) != 1 {
		t.Fatalf("protocolVersion = %v, want 1", field(t, init, "protocolVersion"))
	}

	// authenticate — sent unconditionally with a method id we never offered,
	// and its failure is not caught on the other end.
	c.raw("authenticate", t3Authenticate)

	// session/new — only sessionId is consumed, and it must be a string.
	root := t.TempDir()
	newSession := c.raw("session/new", `{"cwd":`+quote(root)+`,"mcpServers":[]}`)
	id, ok := field(t, newSession, "sessionId").(string)
	if !ok || id == "" {
		t.Fatalf("sessionId = %v", field(t, newSession, "sessionId"))
	}

	// session/prompt — only stopReason is read, and it has to be one of five.
	prompt := c.raw("session/prompt", `{"sessionId":`+quote(id)+`,"prompt":[{"type":"text","text":"say hello"}]}`)
	stop, ok := field(t, prompt, "stopReason").(string)
	if !ok {
		t.Fatalf("stopReason = %v", field(t, prompt, "stopReason"))
	}
	valid := []string{"end_turn", "max_tokens", "max_turn_requests", "refusal", "cancelled"}
	if !slices.Contains(valid, stop) {
		t.Fatalf("stopReason %q is not one the client's schema allows", stop)
	}

	// session/cancel is a notification: no id, nothing to answer.
	c.write(acpMessage{JSONRPC: "2.0", Method: "session/cancel", Params: json.RawMessage(`{"sessionId":` + quote(id) + `}`)})
}

func TestEveryNotificationDecodesAgainstTheClientSchema(t *testing.T) {
	f := newFakeAPI(t,
		msg("tool_use", toolUseBlock("t1", "write_file", `{"path":"note.txt","content":"hi"}`)),
		msg("end_turn", textBlock("done")),
	)
	c := newTestServer(t, f, Options{Permission: ModeFullAccess})

	id := c.handshake(t.TempDir())
	c.prompt(id, "write a note")

	updates := c.sessionUpdates(id)
	if len(updates) == 0 {
		t.Fatal("the turn produced no updates at all")
	}

	// The variants the client switches on. Anything else it drops silently,
	// which is a worse failure than a loud one: the row simply never appears.
	known := []string{
		"agent_message_chunk", "agent_thought_chunk", "tool_call", "tool_call_update",
		"plan", "current_mode_update", "config_option_update", "available_commands_update",
	}

	var sawToolCall, sawToolUpdate bool
	for _, u := range updates {
		if !slices.Contains(known, u.SessionUpdate) {
			t.Errorf("sessionUpdate %q is not one the client handles", u.SessionUpdate)
		}
		switch u.SessionUpdate {
		case "agent_message_chunk", "agent_thought_chunk":
			// The client reads content.type == "text" and content.text.
			raw, _ := json.Marshal(u.Content)
			var content struct {
				Type string `json:"type"`
				Text string `json:"text"`
			}
			if err := json.Unmarshal(raw, &content); err != nil || content.Type != "text" {
				t.Errorf("chunk content is not a text block: %s", raw)
			}
		case "tool_call":
			sawToolCall = true
			// title and toolCallId are the two required fields.
			if u.ToolCallID == "" || u.Title == "" {
				t.Errorf("tool_call is missing a required field: %+v", u)
			}
			if !slices.Contains([]string{"read", "edit", "delete", "move", "search", "execute", "think", "fetch", "switch_mode", "other"}, u.Kind) {
				t.Errorf("tool kind %q is not in the client's vocabulary", u.Kind)
			}
			if !slices.Contains([]string{"pending", "in_progress", "completed", "failed"}, u.Status) {
				t.Errorf("tool status %q is not in the client's vocabulary", u.Status)
			}
		case "tool_call_update":
			sawToolUpdate = true
			if u.ToolCallID == "" {
				t.Error("tool_call_update with no toolCallId cannot be matched to its call")
			}
		}
	}

	if !sawToolCall || !sawToolUpdate {
		t.Errorf("a turn that ran a tool produced no tool rows: %+v", updates)
	}
}

func TestAPermissionRequestMatchesTheClientSchema(t *testing.T) {
	f := newFakeAPI(t,
		msg("tool_use", toolUseBlock("t1", "write_file", `{"path":"note.txt","content":"hi"}`)),
		msg("end_turn", textBlock("done")),
	)
	c := newTestServer(t, f, Options{Permission: ModeAsk})

	c.onPermission = func(acp.RequestPermissionParams) acp.RequestPermissionResult {
		return acp.RequestPermissionResult{
			Outcome: acp.PermissionOutcome{Outcome: acp.OutcomeSelected, OptionID: acp.OptAllowOnce},
		}
	}

	id := c.handshake(t.TempDir())
	c.prompt(id, "write a note")

	c.mu.Lock()
	seen := c.permission
	c.mu.Unlock()
	if seen == nil {
		t.Fatal("no permission request was sent")
	}

	var params struct {
		SessionID string `json:"sessionId"`
		ToolCall  struct {
			ToolCallID string `json:"toolCallId"`
		} `json:"toolCall"`
		Options []struct {
			OptionID string `json:"optionId"`
			Name     string `json:"name"`
			Kind     string `json:"kind"`
		} `json:"options"`
	}
	if err := json.Unmarshal(seen, &params); err != nil {
		t.Fatalf("the request does not decode: %v\n%s", err, seen)
	}
	if params.SessionID != id || params.ToolCall.ToolCallID == "" {
		t.Errorf("params = %s", seen)
	}
	if len(params.Options) == 0 {
		t.Fatal("a permission request with no options cannot be answered")
	}

	// The option kinds are a closed set, and the ids have to include the ones
	// a client may reply with regardless of what we offered.
	var ids []string
	for _, o := range params.Options {
		if o.Name == "" {
			t.Errorf("an option with no name cannot be shown: %+v", o)
		}
		if !slices.Contains([]string{"allow_once", "allow_always", "reject_once", "reject_always"}, o.Kind) {
			t.Errorf("option kind %q is not in the client's vocabulary", o.Kind)
		}
		ids = append(ids, o.OptionID)
	}
	for _, want := range []string{"allow-once", "allow-always", "reject-once"} {
		if !slices.Contains(ids, want) {
			t.Errorf("option id %q is missing; a client that replies with it would not be understood. got %v", want, ids)
		}
	}
}

// quote is strconv.Quote for the handful of JSON strings built by hand here.
func quote(s string) string {
	b, _ := json.Marshal(s)
	return string(b)
}

// stray is a guard: nothing but protocol frames may reach the stream.
func TestNothingButFramesReachTheStream(t *testing.T) {
	f := newFakeAPI(t, msg("end_turn", textBlock("quiet")))
	c := newTestServer(t, f, Options{})

	id := c.handshake(t.TempDir())
	c.prompt(id, "say something")

	for _, line := range c.lines() {
		if strings.TrimSpace(line) == "" {
			continue
		}
		var probe map[string]any
		if err := json.Unmarshal([]byte(line), &probe); err != nil {
			t.Fatalf("a non-JSON line reached the stream and would desynchronise the parser: %q", line)
		}
		if probe["jsonrpc"] != "2.0" {
			t.Errorf("a frame without jsonrpc 2.0: %q", line)
		}
	}
}

func TestTheProbeCanReadTheModelsFromTheHandshakeAlone(t *testing.T) {
	// A health check spawns the agent, shakes hands and stops. T3 Code reads
	// initialize._meta.modelState on exactly that path, so a model list only
	// offered at session/new never reaches the provider card.
	c := newTestServer(t, newFakeAPI(t), Options{
		Models:       []acp.ModelInfo{{ModelID: "codey-opus-5", Name: "Codey Opus 5"}},
		DefaultModel: "codey-opus-5",
	})

	init := c.raw("initialize", t3Initialize)

	var resp struct {
		Meta struct {
			ModelState struct {
				CurrentModelID  string `json:"currentModelId"`
				AvailableModels []struct {
					ModelID string `json:"modelId"`
					Name    string `json:"name"`
				} `json:"availableModels"`
			} `json:"modelState"`
		} `json:"_meta"`
	}
	if err := json.Unmarshal(init, &resp); err != nil {
		t.Fatalf("initialize result: %v", err)
	}

	state := resp.Meta.ModelState
	if state.CurrentModelID != "codey-opus-5" || len(state.AvailableModels) != 1 {
		t.Fatalf("_meta.modelState = %+v, from %s", state, init)
	}
	if state.AvailableModels[0].Name != "Codey Opus 5" {
		t.Errorf("model name = %q", state.AvailableModels[0].Name)
	}
}