feat(control-plane): AG-UI protocol adapter (POC)

control-plane/internal/agui/events_test.go

@@ -0,0 +1,266 @@
+package agui
+
+import (
+	"bytes"
+	"encoding/json"
+	"strings"
+	"testing"
+)
+
+// TestWriteSSE_FrameShape pins the canonical AG-UI wire format:
+//   - frame is `data: <json>\n\n` only (no `event:` line — see encoder.ts /
+//     encoder.py in ag-ui-protocol/ag-ui)
+//   - `type` field carries the UPPER_SNAKE_CASE event name
+//   - timestamp, when present, is a number (Unix ms)
+func TestWriteSSE_FrameShape(t *testing.T) {
+	cases := []struct {
+		name       string
+		ev         Event
+		wantTyp    string
+		wantFields []string
+	}{
+		{
+			name:       "RunStarted",
+			ev:         RunStarted{ThreadID: "thread-1", RunID: "run-1", Timestamp: 1700000000000},
+			wantTyp:    "RUN_STARTED",
+			wantFields: []string{`"threadId":"thread-1"`, `"runId":"run-1"`, `"timestamp":1700000000000`},
+		},
+		{
+			name:       "RunFinished_success_carriesIDs",
+			ev:         RunFinished{ThreadID: "thread-1", RunID: "run-1", Outcome: &Outcome{Type: "success"}, Result: map[string]any{"answer": 42}},
+			wantTyp:    "RUN_FINISHED",
+			wantFields: []string{`"threadId":"thread-1"`, `"runId":"run-1"`, `"outcome":{"type":"success"}`, `"answer":42`},
+		},
+		{
+			name:       "RunError",
+			ev:         RunError{Message: "boom", Code: "ERR_X"},
+			wantTyp:    "RUN_ERROR",
+			wantFields: []string{`"message":"boom"`, `"code":"ERR_X"`},
+		},
+		{
+			name:       "TextMessageStart",
+			ev:         TextMessageStart{MessageID: "msg-1", Role: "assistant"},
+			wantTyp:    "TEXT_MESSAGE_START",
+			wantFields: []string{`"messageId":"msg-1"`, `"role":"assistant"`},
+		},
+		{
+			name:       "TextMessageContent",
+			ev:         TextMessageContent{MessageID: "msg-1", Delta: "hello"},
+			wantTyp:    "TEXT_MESSAGE_CONTENT",
+			wantFields: []string{`"messageId":"msg-1"`, `"delta":"hello"`},
+		},
+		{
+			name:       "TextMessageEnd",
+			ev:         TextMessageEnd{MessageID: "msg-1"},
+			wantTyp:    "TEXT_MESSAGE_END",
+			wantFields: []string{`"messageId":"msg-1"`},
+		},
+		{
+			name:       "ToolCallStart",
+			ev:         ToolCallStart{ToolCallID: "tc-1", ToolCallName: "showFlightCard", ParentMessageID: "msg-1"},
+			wantTyp:    "TOOL_CALL_START",
+			wantFields: []string{`"toolCallId":"tc-1"`, `"toolCallName":"showFlightCard"`, `"parentMessageId":"msg-1"`},
+		},
+		{
+			name:       "ToolCallArgs",
+			ev:         ToolCallArgs{ToolCallID: "tc-1", Delta: `{"from":"SFO"}`},
+			wantTyp:    "TOOL_CALL_ARGS",
+			wantFields: []string{`"toolCallId":"tc-1"`, `"delta":"{\"from\":\"SFO\"}"`},
+		},
+		{
+			name:       "ToolCallEnd",
+			ev:         ToolCallEnd{ToolCallID: "tc-1"},
+			wantTyp:    "TOOL_CALL_END",
+			wantFields: []string{`"toolCallId":"tc-1"`},
+		},
+		{
+			name:       "ToolCallResult",
+			ev:         ToolCallResult{MessageID: "msg-2", ToolCallID: "tc-1", Content: "ok", Role: "tool"},
+			wantTyp:    "TOOL_CALL_RESULT",
+			wantFields: []string{`"messageId":"msg-2"`, `"toolCallId":"tc-1"`, `"content":"ok"`, `"role":"tool"`},
+		},
+		{
+			name:       "MessagesSnapshot",
+			ev:         MessagesSnapshot{Messages: []Message{{ID: "m1", Role: "user", Content: "hi"}}},
+			wantTyp:    "MESSAGES_SNAPSHOT",
+			wantFields: []string{`"messages":[`, `"role":"user"`, `"content":"hi"`},
+		},
+		{
+			name:       "StateSnapshot",
+			ev:         StateSnapshot{Snapshot: map[string]any{"counter": 1}},
+			wantTyp:    "STATE_SNAPSHOT",
+			wantFields: []string{`"snapshot":{"counter":1}`},
+		},
+		{
+			name:       "StateDelta",
+			ev:         StateDelta{Delta: []any{map[string]any{"op": "replace", "path": "/counter", "value": 2}}},
+			wantTyp:    "STATE_DELTA",
+			wantFields: []string{`"delta":[`, `"op":"replace"`, `"path":"/counter"`},
+		},
+		{
+			name:       "StepStarted",
+			ev:         StepStarted{StepName: "plan"},
+			wantTyp:    "STEP_STARTED",
+			wantFields: []string{`"stepName":"plan"`},
+		},
+		{
+			name:       "StepFinished",
+			ev:         StepFinished{StepName: "plan"},
+			wantTyp:    "STEP_FINISHED",
+			wantFields: []string{`"stepName":"plan"`},
+		},
+		{
+			name:       "RawEvent",
+			ev:         RawEvent{Event: map[string]any{"foo": 1}, Source: "harness"},
+			wantTyp:    "RAW",
+			wantFields: []string{`"event":{"foo":1}`, `"source":"harness"`},
+		},
+		{
+			name:       "CustomEvent",
+			ev:         CustomEvent{Name: "ack", Value: map[string]any{"ok": true}},
+			wantTyp:    "CUSTOM",
+			wantFields: []string{`"name":"ack"`, `"value":{"ok":true}`},
+		},
+		{
+			name:       "ReasoningStart",
+			ev:         ReasoningStart{MessageID: "r1"},
+			wantTyp:    "REASONING_START",
+			wantFields: []string{`"messageId":"r1"`},
+		},
+		{
+			name:       "ReasoningMessageStart",
+			ev:         ReasoningMessageStart{MessageID: "r1", Role: "reasoning"},
+			wantTyp:    "REASONING_MESSAGE_START",
+			wantFields: []string{`"messageId":"r1"`, `"role":"reasoning"`},
+		},
+		{
+			name:       "ReasoningMessageContent",
+			ev:         ReasoningMessageContent{MessageID: "r1", Delta: "thinking..."},
+			wantTyp:    "REASONING_MESSAGE_CONTENT",
+			wantFields: []string{`"messageId":"r1"`, `"delta":"thinking..."`},
+		},
+		{
+			name:       "ReasoningMessageEnd",
+			ev:         ReasoningMessageEnd{MessageID: "r1"},
+			wantTyp:    "REASONING_MESSAGE_END",
+			wantFields: []string{`"messageId":"r1"`},
+		},
+		{
+			name:       "ReasoningEnd",
+			ev:         ReasoningEnd{MessageID: "r1"},
+			wantTyp:    "REASONING_END",
+			wantFields: []string{`"messageId":"r1"`},
+		},
+		{
+			name:       "TextMessageChunk",
+			ev:         TextMessageChunk{MessageID: "m1", Delta: "tok"},
+			wantTyp:    "TEXT_MESSAGE_CHUNK",
+			wantFields: []string{`"messageId":"m1"`, `"delta":"tok"`},
+		},
+		{
+			name:       "ToolCallChunk",
+			ev:         ToolCallChunk{ToolCallID: "tc1", ToolCallName: "go", Delta: "{\"a\":1}"},
+			wantTyp:    "TOOL_CALL_CHUNK",
+			wantFields: []string{`"toolCallId":"tc1"`, `"toolCallName":"go"`, `"delta":"{\"a\":1}"`},
+		},
+	}
+
+	for _, tc := range cases {
+		t.Run(tc.name, func(t *testing.T) {
+			var buf bytes.Buffer
+			if err := WriteSSE(&buf, tc.ev); err != nil {
+				t.Fatalf("WriteSSE: %v", err)
+			}
+			frame := buf.String()
+
+			// Canonical wire shape: `data: <json>\n\n`. No `event:` line.
+			if !strings.HasPrefix(frame, "data: ") {
+				t.Fatalf("frame must start with `data: `:\n%s", frame)
+			}
+			if !strings.HasSuffix(frame, "\n\n") {
+				t.Fatalf("frame must end with blank-line terminator:\n%s", frame)
+			}
+			if strings.Contains(frame, "\nevent:") || strings.HasPrefix(frame, "event:") {
+				t.Fatalf("frame must not include an `event:` line (canonical encoder omits it):\n%s", frame)
+			}
+
+			body := strings.TrimSuffix(strings.TrimPrefix(frame, "data: "), "\n\n")
+			var decoded map[string]any
+			if err := json.Unmarshal([]byte(body), &decoded); err != nil {
+				t.Fatalf("data line is not JSON: %v\nbody: %s", err, body)
+			}
+			if got := decoded["type"]; got != tc.wantTyp {
+				t.Fatalf("json type field = %v, want %q", got, tc.wantTyp)
+			}
+			for _, want := range tc.wantFields {
+				if !strings.Contains(body, want) {
+					t.Fatalf("expected field %s in payload:\n%s", want, body)
+				}
+			}
+		})
+	}
+}
+
+// TestWriteSSE_OmitsZeroOptionalFields confirms our `omitempty` tags drop
+// timestamp / role / outcome / code when they're at zero values, matching
+// the Python encoder's `exclude_none=True` semantics.
+func TestWriteSSE_OmitsZeroOptionalFields(t *testing.T) {
+	var buf bytes.Buffer
+	if err := WriteSSE(&buf, TextMessageStart{MessageID: "m"}); err != nil {
+		t.Fatal(err)
+	}
+	body := buf.String()
+	if strings.Contains(body, `"role":""`) {
+		t.Errorf("empty role should be omitted: %s", body)
+	}
+	if strings.Contains(body, `"timestamp":0`) {
+		t.Errorf("zero timestamp should be omitted: %s", body)
+	}
+}
+
+// unmarshalableEvent fails JSON encoding deterministically so we can exercise
+// the marshal-error branch in WriteSSE.
+type unmarshalableEvent struct{}
+
+func (unmarshalableEvent) Type() string                 { return "BAD_EVENT" }
+func (unmarshalableEvent) MarshalJSON() ([]byte, error) { return nil, errBoom }
+
+var errBoom = &boomError{}
+
+type boomError struct{}
+
+func (b *boomError) Error() string { return "boom" }
+
+// TestWriteSSE_MarshalErrorIsReturned ensures encode failures surface to the
+// caller rather than producing a silently-malformed frame.
+func TestWriteSSE_MarshalErrorIsReturned(t *testing.T) {
+	var buf bytes.Buffer
+	err := WriteSSE(&buf, unmarshalableEvent{})
+	if err == nil {
+		t.Fatalf("expected marshal error, got nil; buf=%q", buf.String())
+	}
+	if !strings.Contains(err.Error(), "marshal BAD_EVENT") {
+		t.Errorf("error should name the event type: %v", err)
+	}
+	if buf.Len() != 0 {
+		t.Errorf("nothing should be written on marshal failure; got %q", buf.String())
+	}
+}
+
+// failingWriter returns an error on every Write — used to cover the
+// write-error branch of WriteSSE.
+type failingWriter struct{}
+
+func (failingWriter) Write([]byte) (int, error) { return 0, errBoom }
+
+// TestWriteSSE_WriteErrorIsReturned confirms a flaky writer surfaces to the
+// caller (the handler uses this to bail out cleanly on client disconnect).
+func TestWriteSSE_WriteErrorIsReturned(t *testing.T) {
+	err := WriteSSE(failingWriter{}, RunStarted{ThreadID: "t", RunID: "r"})
+	if err == nil {
+		t.Fatalf("expected write error, got nil")
+	}
+	if !strings.Contains(err.Error(), "write RUN_STARTED") {
+		t.Errorf("error should name the event type: %v", err)
+	}
+}

control-plane/internal/agui/types.go

@@ -0,0 +1,77 @@
+package agui
+
+import "encoding/json"
+
+// RunAgentInput mirrors the canonical RunAgentInputSchema from the AG-UI
+// reference SDK (sdks/typescript/packages/core/src/types.ts). The vanilla
+// @ag-ui/client HttpAgent — and the CopilotRuntime that wraps it — POSTs
+// this exact shape to backends.
+//
+// We keep all fields permissive (json.RawMessage / any) so unrecognized
+// or evolving sub-fields pass through without forcing schema bumps.
+type RunAgentInput struct {
+	ThreadID       string            `json:"threadId"`
+	RunID          string            `json:"runId"`
+	ParentRunID    string            `json:"parentRunId,omitempty"`
+	State          json.RawMessage   `json:"state,omitempty"`
+	Messages       []Message         `json:"messages,omitempty"`
+	Tools          []Tool            `json:"tools,omitempty"`
+	Context        []ContextItem     `json:"context,omitempty"`
+	ForwardedProps json.RawMessage   `json:"forwardedProps,omitempty"`
+	Resume         []json.RawMessage `json:"resume,omitempty"`
+}
+
+// Message is the canonical AG-UI message envelope (MessageSchema). Role
+// drives the discriminated union; we keep optional fields so user/assistant/
+// tool messages all round-trip through the same struct.
+type Message struct {
+	ID         string     `json:"id,omitempty"`
+	Role       string     `json:"role"`
+	Content    string     `json:"content,omitempty"`
+	Name       string     `json:"name,omitempty"`
+	ToolCallID string     `json:"toolCallId,omitempty"`
+	ToolCalls  []ToolCall `json:"toolCalls,omitempty"`
+}
+
+// ToolCall is an assistant-message-attached tool invocation, matching
+// ToolCallSchema. The function arguments are a JSON string per OpenAI
+// convention.
+type ToolCall struct {
+	ID       string           `json:"id"`
+	Type     string           `json:"type"` // always "function" today
+	Function ToolCallFunction `json:"function"`
+}
+
+type ToolCallFunction struct {
+	Name      string `json:"name"`
+	Arguments string `json:"arguments"`
+}
+
+// Tool describes a tool the frontend has registered (e.g. via
+// useCopilotAction). Reasoners can choose to invoke these by emitting a
+// matching TOOL_CALL_* sequence.
+type Tool struct {
+	Name        string          `json:"name"`
+	Description string          `json:"description,omitempty"`
+	Parameters  json.RawMessage `json:"parameters,omitempty"` // JSON Schema
+	Metadata    json.RawMessage `json:"metadata,omitempty"`
+}
+
+// ContextItem is one (description, value) pair from the readables stream
+// (e.g. useCopilotReadable). Value is freeform JSON.
+type ContextItem struct {
+	Description string          `json:"description,omitempty"`
+	Value       json.RawMessage `json:"value,omitempty"`
+}
+
+// LastUserMessageText returns the trailing user-role message's content,
+// which is the conventional "prompt" for chat-style agents. Empty string
+// if the trailing message is not user-role or messages is empty.
+func (r RunAgentInput) LastUserMessageText() string {
+	for i := len(r.Messages) - 1; i >= 0; i-- {
+		if r.Messages[i].Role == "user" {
+			return r.Messages[i].Content
+		}
+	}
+	return ""
+}