feat(agentx): add OpenAI-compatible streaming LLM client

Parse MiniMax-style SSE deltas, including incremental tool calls, and encode transport failures on the event stream instead of panicking. Co-authored-by: Cursor <cursoragent@cursor.com>
2026-09-21 14:43:32 +00:00
parent 6dc5f1cfa6
commit b49b95dc7a
5 changed files with 699 additions and 0 deletions
@@ -0,0 +1,152 @@
+package llm
+
+import (
+	"bytes"
+	"context"
+	"encoding/json"
+	"fmt"
+	"io"
+	"net/http"
+	"strings"
+	"time"
+)
+
+// Client calls an OpenAI-compatible chat completions endpoint with streaming.
+type Client struct {
+	BaseURL    string
+	APIKey     string
+	HTTPClient *http.Client
+}
+
+// Stream implements StreamFunc. Transport and provider failures are encoded as
+// KindError events; the returned channel is always non-nil and is closed.
+func (c *Client) Stream(ctx context.Context, model string, messages []Message, tools []ToolDecl, opts StreamOptions) <-chan Event {
+	out := make(chan Event)
+	go func() {
+		defer close(out)
+		c.run(ctx, model, messages, tools, opts, out)
+	}()
+	return out
+}
+
+func (c *Client) run(ctx context.Context, model string, messages []Message, tools []ToolDecl, opts StreamOptions, out chan<- Event) {
+	base := strings.TrimRight(c.BaseURL, "/")
+	if opts.APIBase != "" {
+		base = strings.TrimRight(opts.APIBase, "/")
+	}
+	key := c.APIKey
+	if opts.APIKey != "" {
+		key = opts.APIKey
+	}
+	if base == "" || key == "" || model == "" {
+		emitError(out, "", false, "agentx llm client is missing base url, api key, or model")
+		return
+	}
+	body := map[string]any{
+		"model":    model,
+		"messages": toProviderMessages(messages),
+		"stream":   true,
+		"stream_options": map[string]any{
+			"include_usage": true,
+		},
+	}
+	if opts.MaxTokens > 0 {
+		body["max_tokens"] = opts.MaxTokens
+	}
+	if len(tools) > 0 {
+		body["tools"] = toProviderTools(tools)
+		body["tool_choice"] = "auto"
+	}
+	raw, err := json.Marshal(body)
+	if err != nil {
+		emitError(out, "", false, err.Error())
+		return
+	}
+	req, err := http.NewRequestWithContext(ctx, http.MethodPost, base+"/chat/completions", bytes.NewReader(raw))
+	if err != nil {
+		emitError(out, "", false, err.Error())
+		return
+	}
+	req.Header.Set("Authorization", "Bearer "+key)
+	req.Header.Set("Content-Type", "application/json")
+	req.Header.Set("Accept", "text/event-stream")
+	hc := c.HTTPClient
+	if hc == nil {
+		hc = &http.Client{Timeout: 30 * time.Minute}
+	}
+	resp, err := hc.Do(req)
+	if err != nil {
+		if ctx.Err() != nil {
+			emitAbort(out, "", ctx.Err())
+			return
+		}
+		emitError(out, "", true, err.Error())
+		return
+	}
+	defer resp.Body.Close()
+	if resp.StatusCode >= 400 {
+		slurp, _ := io.ReadAll(io.LimitReader(resp.Body, 4096))
+		retryable := resp.StatusCode == http.StatusTooManyRequests || resp.StatusCode >= 500
+		emitError(out, "", retryable, fmt.Sprintf("provider status %d: %s", resp.StatusCode, strings.TrimSpace(string(slurp))))
+		return
+	}
+	for ev := range ParseSSE(ctx, resp.Body) {
+		out <- ev
+	}
+}
+
+func toProviderMessages(messages []Message) []map[string]any {
+	out := make([]map[string]any, 0, len(messages))
+	for _, m := range messages {
+		switch m.Role {
+		case RoleTool:
+			out = append(out, map[string]any{
+				"role":         "tool",
+				"tool_call_id": m.ToolCallID,
+				"content":      m.Content,
+			})
+		case RoleAssistant:
+			item := map[string]any{"role": "assistant", "content": m.Content}
+			if len(m.ToolCalls) > 0 {
+				calls := make([]map[string]any, 0, len(m.ToolCalls))
+				for _, tc := range m.ToolCalls {
+					calls = append(calls, map[string]any{
+						"id":   tc.ID,
+						"type": "function",
+						"function": map[string]any{
+							"name":      tc.Name,
+							"arguments": tc.Arguments,
+						},
+					})
+				}
+				item["tool_calls"] = calls
+			}
+			out = append(out, item)
+		default:
+			out = append(out, map[string]any{
+				"role":    string(m.Role),
+				"content": m.Content,
+			})
+		}
+	}
+	return out
+}
+
+func toProviderTools(tools []ToolDecl) []map[string]any {
+	out := make([]map[string]any, 0, len(tools))
+	for _, t := range tools {
+		params := t.Parameters
+		if len(params) == 0 {
+			params = json.RawMessage(`{"type":"object","properties":{}}`)
+		}
+		out = append(out, map[string]any{
+			"type": "function",
+			"function": map[string]any{
+				"name":        t.Name,
+				"description": t.Description,
+				"parameters":  json.RawMessage(params),
+			},
+		})
+	}
+	return out
+}
@@ -0,0 +1,50 @@
+package llm
+
+import (
+	"context"
+	"os"
+	"testing"
+	"time"
+)
+
+// TestSmokeMiniMax hits the real MiniMax endpoint when RUN_AGENTX_SMOKE=1.
+// It calibrates the incremental SSE tool-call format and is not part of -short CI.
+func TestSmokeMiniMax(t *testing.T) {
+	if os.Getenv("RUN_AGENTX_SMOKE") != "1" {
+		t.Skip("set RUN_AGENTX_SMOKE=1 to call MiniMax")
+	}
+	key := os.Getenv("AGENTX_API_KEY")
+	if key == "" {
+		key = os.Getenv("MINIMAX_API_KEY")
+	}
+	if key == "" {
+		t.Fatal("AGENTX_API_KEY or MINIMAX_API_KEY is required")
+	}
+	base := os.Getenv("AGENTX_API_BASE")
+	if base == "" {
+		base = os.Getenv("MINIMAX_ENDPOINT_URL")
+	}
+	if base == "" {
+		base = "https://api.minimaxi.com/v1"
+	}
+	model := os.Getenv("AGENTX_MODEL")
+	if model == "" {
+		model = "MiniMax-M2.7-highspeed"
+	}
+	ctx, cancel := context.WithTimeout(context.Background(), 90*time.Second)
+	defer cancel()
+	c := &Client{BaseURL: base, APIKey: key, HTTPClient: nil}
+	var end Event
+	for ev := range c.Stream(ctx, model, []Message{{Role: RoleUser, Content: "Reply with the single word pong."}}, nil, StreamOptions{MaxTokens: 64}) {
+		if ev.Kind == KindError {
+			t.Fatalf("provider error retryable=%v %s", ev.Retryable, ev.ErrorMessage)
+		}
+		if ev.Kind == KindMessageEnd {
+			end = ev
+		}
+	}
+	if end.Message.StopReason != StopEnd || end.Message.Content == "" {
+		t.Fatalf("unexpected completion %#v", end.Message)
+	}
+	t.Logf("smoke ok stop=%s tokens=%d", end.Message.StopReason, end.Usage.TotalTokens)
+}
@@ -0,0 +1,211 @@
+package llm
+
+import (
+	"bufio"
+	"context"
+	"encoding/json"
+	"io"
+	"strings"
+)
+
+type sseChoice struct {
+	Delta struct {
+		Content   string `json:"content"`
+		ToolCalls []struct {
+			Index    int    `json:"index"`
+			ID       string `json:"id"`
+			Function struct {
+				Name      string `json:"name"`
+				Arguments string `json:"arguments"`
+			} `json:"function"`
+		} `json:"tool_calls"`
+	} `json:"delta"`
+	FinishReason string `json:"finish_reason"`
+}
+
+type sseChunk struct {
+	Choices []sseChoice `json:"choices"`
+	Usage   *Usage      `json:"usage"`
+	Error   *struct {
+		Message string `json:"message"`
+	} `json:"error"`
+}
+
+// ParseSSE reads an OpenAI-compatible server-sent event body and emits
+// deltas plus one terminal message_end or error event.
+func ParseSSE(ctx context.Context, r io.Reader) <-chan Event {
+	out := make(chan Event)
+	go func() {
+		defer close(out)
+		emitParsed(ctx, r, out)
+	}()
+	return out
+}
+
+func emitParsed(ctx context.Context, r io.Reader, out chan<- Event) {
+	var content strings.Builder
+	calls := map[int]*ToolCall{}
+	order := []int{}
+	var usage Usage
+	finished := ""
+	sc := bufio.NewScanner(r)
+	sc.Buffer(make([]byte, 0, 64*1024), 4*1024*1024)
+	for sc.Scan() {
+		if ctx.Err() != nil {
+			emitAbort(out, content.String(), ctx.Err())
+			return
+		}
+		line := strings.TrimSpace(sc.Text())
+		if line == "" || strings.HasPrefix(line, ":") {
+			continue
+		}
+		if !strings.HasPrefix(line, "data:") {
+			continue
+		}
+		data := strings.TrimSpace(strings.TrimPrefix(line, "data:"))
+		if data == "[DONE]" {
+			break
+		}
+		var chunk sseChunk
+		if err := json.Unmarshal([]byte(data), &chunk); err != nil {
+			emitError(out, content.String(), false, "malformed stream chunk: "+err.Error())
+			return
+		}
+		if chunk.Error != nil && chunk.Error.Message != "" {
+			emitError(out, content.String(), false, chunk.Error.Message)
+			return
+		}
+		if chunk.Usage != nil {
+			usage = *chunk.Usage
+		}
+		if len(chunk.Choices) == 0 {
+			continue
+		}
+		choice := chunk.Choices[0]
+		if choice.FinishReason != "" {
+			finished = choice.FinishReason
+		}
+		if choice.Delta.Content != "" {
+			content.WriteString(choice.Delta.Content)
+			select {
+			case <-ctx.Done():
+				emitAbort(out, content.String(), ctx.Err())
+				return
+			case out <- Event{Kind: KindTextDelta, Text: choice.Delta.Content}:
+			}
+		}
+		for _, tc := range choice.Delta.ToolCalls {
+			cur, ok := calls[tc.Index]
+			if !ok {
+				cur = &ToolCall{}
+				calls[tc.Index] = cur
+				order = append(order, tc.Index)
+			}
+			if tc.ID != "" {
+				cur.ID = tc.ID
+			}
+			if tc.Function.Name != "" {
+				cur.Name += tc.Function.Name
+			}
+			if tc.Function.Arguments != "" {
+				cur.Arguments += tc.Function.Arguments
+			}
+			select {
+			case <-ctx.Done():
+				emitAbort(out, content.String(), ctx.Err())
+				return
+			case out <- Event{
+				Kind:           KindToolCallDelta,
+				ToolIndex:      tc.Index,
+				ToolID:         tc.ID,
+				ToolName:       tc.Function.Name,
+				ArgumentsDelta: tc.Function.Arguments,
+			}:
+			}
+		}
+	}
+	if err := sc.Err(); err != nil {
+		if ctx.Err() != nil {
+			emitAbort(out, content.String(), ctx.Err())
+			return
+		}
+		emitError(out, content.String(), true, err.Error())
+		return
+	}
+	msg := Message{Role: RoleAssistant, Content: content.String(), StopReason: mapFinish(finished, len(order) > 0)}
+	for _, idx := range order {
+		msg.ToolCalls = append(msg.ToolCalls, *calls[idx])
+	}
+	if msg.StopReason == StopToolCalls && len(msg.ToolCalls) == 0 {
+		msg.StopReason = StopEnd
+	}
+	select {
+	case <-ctx.Done():
+		emitAbort(out, content.String(), ctx.Err())
+	case out <- Event{Kind: KindMessageEnd, Message: msg, Usage: usage}:
+	}
+}
+
+func mapFinish(reason string, hasTools bool) StopReason {
+	switch reason {
+	case "tool_calls", "function_call":
+		return StopToolCalls
+	case "content_filter", "error":
+		return StopError
+	case "stop", "length", "":
+		if hasTools {
+			return StopToolCalls
+		}
+		return StopEnd
+	default:
+		if hasTools {
+			return StopToolCalls
+		}
+		return StopEnd
+	}
+}
+
+func emitError(out chan<- Event, partial string, retryable bool, message string) {
+	out <- Event{
+		Kind:         KindError,
+		Retryable:    retryable,
+		ErrorMessage: message,
+		Err:          &StreamError{Message: message, Retryable: retryable},
+		Message: Message{
+			Role:       RoleAssistant,
+			Content:    partial,
+			StopReason: StopError,
+		},
+	}
+}
+
+func emitAbort(out chan<- Event, partial string, err error) {
+	msg := "aborted"
+	if err != nil {
+		msg = err.Error()
+	}
+	out <- Event{
+		Kind:         KindError,
+		Retryable:    false,
+		ErrorMessage: msg,
+		Err:          err,
+		Message: Message{
+			Role:       RoleAssistant,
+			Content:    partial,
+			StopReason: StopAborted,
+		},
+	}
+}
+
+// StreamError is a provider or transport failure encoded on the event stream.
+type StreamError struct {
+	Message   string
+	Retryable bool
+}
+
+func (e *StreamError) Error() string {
+	if e == nil {
+		return "stream error"
+	}
+	return e.Message
+}
@@ -0,0 +1,191 @@
+package llm
+
+import (
+	"context"
+	"io"
+	"net/http"
+	"net/http/httptest"
+	"strings"
+	"testing"
+	"time"
+)
+
+func TestParseSSETextAndUsage(t *testing.T) {
+	body := strings.Join([]string{
+		`data: {"choices":[{"delta":{"content":"你"},"index":0}]}`,
+		`data: {"choices":[{"delta":{"content":"好"},"finish_reason":"stop"}],"usage":{"prompt_tokens":3,"completion_tokens":2,"total_tokens":5}}`,
+		`data: [DONE]`,
+		"",
+	}, "\n")
+	var deltas []string
+	var end Event
+	for ev := range ParseSSE(context.Background(), strings.NewReader(body)) {
+		switch ev.Kind {
+		case KindTextDelta:
+			deltas = append(deltas, ev.Text)
+		case KindMessageEnd:
+			end = ev
+		default:
+			t.Fatalf("unexpected event %#v", ev)
+		}
+	}
+	if strings.Join(deltas, "") != "你好" {
+		t.Fatalf("deltas %v", deltas)
+	}
+	if end.Message.StopReason != StopEnd || end.Message.Content != "你好" {
+		t.Fatalf("end %#v", end.Message)
+	}
+	if end.Usage.TotalTokens != 5 {
+		t.Fatalf("usage %#v", end.Usage)
+	}
+}
+
+func TestParseSSEToolCallIncrements(t *testing.T) {
+	body := strings.Join([]string{
+		`data: {"choices":[{"delta":{"tool_calls":[{"index":0,"id":"call_1","function":{"name":"read_file","arguments":""}}]}}]}`,
+		`data: {"choices":[{"delta":{"tool_calls":[{"index":0,"function":{"arguments":"{\"path\":"}}]}}]}`,
+		`data: {"choices":[{"delta":{"tool_calls":[{"index":0,"function":{"arguments":"\"a.md\"}"}}]},"finish_reason":"tool_calls"}]}`,
+		`data: [DONE]`,
+		"",
+	}, "\n")
+	var end Event
+	deltas := 0
+	for ev := range ParseSSE(context.Background(), strings.NewReader(body)) {
+		if ev.Kind == KindToolCallDelta {
+			deltas++
+		}
+		if ev.Kind == KindMessageEnd {
+			end = ev
+		}
+	}
+	if deltas != 3 {
+		t.Fatalf("tool deltas %d", deltas)
+	}
+	if end.Message.StopReason != StopToolCalls || len(end.Message.ToolCalls) != 1 {
+		t.Fatalf("end %#v", end.Message)
+	}
+	tc := end.Message.ToolCalls[0]
+	if tc.ID != "call_1" || tc.Name != "read_file" || tc.Arguments != `{"path":"a.md"}` {
+		t.Fatalf("tool call %#v", tc)
+	}
+}
+
+func TestParseSSEMalformedIsError(t *testing.T) {
+	body := "data: {not-json}\n\n"
+	var got Event
+	for ev := range ParseSSE(context.Background(), strings.NewReader(body)) {
+		got = ev
+	}
+	if got.Kind != KindError || got.Message.StopReason != StopError || got.Retryable {
+		t.Fatalf("event %#v", got)
+	}
+}
+
+func TestParseSSECancel(t *testing.T) {
+	pr, pw := io.Pipe()
+	ctx, cancel := context.WithCancel(context.Background())
+	ch := ParseSSE(ctx, pr)
+	go func() {
+		_, _ = io.WriteString(pw, "data: {\"choices\":[{\"delta\":{\"content\":\"x\"}}]}\n\n")
+		cancel()
+		time.Sleep(20 * time.Millisecond)
+		_ = pw.Close()
+	}()
+	sawAbort := false
+	for ev := range ch {
+		if ev.Kind == KindError && ev.Message.StopReason == StopAborted {
+			sawAbort = true
+		}
+	}
+	if !sawAbort {
+		t.Fatal("expected aborted event")
+	}
+}
+
+func TestClientHTTPErrorIsRetryable(t *testing.T) {
+	srv := httpServer(t, func(w http.ResponseWriter, r *http.Request) {
+		if r.Header.Get("Authorization") != "Bearer secret" {
+			t.Errorf("auth header %q", r.Header.Get("Authorization"))
+		}
+		http.Error(w, "busy", http.StatusTooManyRequests)
+	})
+	c := &Client{BaseURL: srv, APIKey: "secret", HTTPClient: srvClient()}
+	var got Event
+	for ev := range c.Stream(context.Background(), "MiniMax-M2.7-highspeed", []Message{{Role: RoleUser, Content: "hi"}}, nil, StreamOptions{}) {
+		got = ev
+	}
+	if got.Kind != KindError || !got.Retryable || got.Message.StopReason != StopError {
+		t.Fatalf("event %#v", got)
+	}
+}
+
+func TestClientStreamsCompletion(t *testing.T) {
+	srv := httpServer(t, func(w http.ResponseWriter, r *http.Request) {
+		w.Header().Set("Content-Type", "text/event-stream")
+		_, _ = io.WriteString(w, "data: {\"choices\":[{\"delta\":{\"content\":\"ok\"},\"finish_reason\":\"stop\"}]}\n\ndata: [DONE]\n")
+	})
+	c := &Client{BaseURL: srv, APIKey: "k", HTTPClient: srvClient()}
+	var end Event
+	for ev := range c.Stream(context.Background(), "m", []Message{{Role: RoleUser, Content: "hi"}}, []ToolDecl{{Name: "read_file", Parameters: []byte(`{"type":"object"}`)}}, StreamOptions{}) {
+		if ev.Kind == KindMessageEnd {
+			end = ev
+		}
+	}
+	if end.Message.Content != "ok" || end.Message.StopReason != StopEnd {
+		t.Fatalf("end %#v", end.Message)
+	}
+}
+
+func TestClientCancel(t *testing.T) {
+	started := make(chan struct{})
+	srv := httpServer(t, func(w http.ResponseWriter, r *http.Request) {
+		close(started)
+		w.Header().Set("Content-Type", "text/event-stream")
+		fl, _ := w.(http.Flusher)
+		_, _ = io.WriteString(w, "data: {\"choices\":[{\"delta\":{\"content\":\"a\"}}]}\n\n")
+		if fl != nil {
+			fl.Flush()
+		}
+		<-r.Context().Done()
+	})
+	c := &Client{BaseURL: srv, APIKey: "k", HTTPClient: &http.Client{}}
+	ctx, cancel := context.WithCancel(context.Background())
+	ch := c.Stream(ctx, "m", []Message{{Role: RoleUser, Content: "hi"}}, nil, StreamOptions{})
+	<-started
+	cancel()
+	saw := false
+	deadline := time.After(2 * time.Second)
+	for {
+		select {
+		case ev, ok := <-ch:
+			if !ok {
+				if !saw {
+					t.Fatal("channel closed without abort")
+				}
+				return
+			}
+			if ev.Kind == KindError && ev.Message.StopReason == StopAborted {
+				saw = true
+			}
+		case <-deadline:
+			t.Fatal("timed out waiting for cancel")
+		}
+	}
+}
+
+func httpServer(t *testing.T, h http.HandlerFunc) string {
+	t.Helper()
+	srv := httptest.NewServer(h)
+	t.Cleanup(srv.Close)
+	return srv.URL
+}
+
+func srvClient() *http.Client {
+	return &http.Client{Timeout: 2 * time.Second}
+}
+
+func TestStreamErrorString(t *testing.T) {
+	if (&StreamError{Message: "boom"}).Error() != "boom" {
+		t.Fatal("error string")
+	}
+}
@@ -0,0 +1,95 @@
+package llm
+
+import (
+	"context"
+	"encoding/json"
+)
+
+// Role is a transcript role carried into the provider request.
+type Role string
+
+const (
+	RoleSystem    Role = "system"
+	RoleUser      Role = "user"
+	RoleAssistant Role = "assistant"
+	RoleTool      Role = "tool"
+)
+
+// StopReason is the terminal condition of one assistant turn.
+type StopReason string
+
+const (
+	StopEnd       StopReason = "end"
+	StopToolCalls StopReason = "toolCalls"
+	StopError     StopReason = "error"
+	StopAborted   StopReason = "aborted"
+)
+
+// EventKind identifies one item on the assistant stream.
+type EventKind string
+
+const (
+	KindTextDelta     EventKind = "text_delta"
+	KindToolCallDelta EventKind = "tool_call_delta"
+	KindMessageEnd    EventKind = "message_end"
+	KindError         EventKind = "error"
+)
+
+// ToolCall is one function call requested by the assistant.
+type ToolCall struct {
+	ID        string `json:"id"`
+	Name      string `json:"name"`
+	Arguments string `json:"arguments"`
+}
+
+// Message is one transcript entry. System, user, assistant, and tool results share this shape.
+type Message struct {
+	Role       Role       `json:"role"`
+	Content    string     `json:"content,omitempty"`
+	Name       string     `json:"name,omitempty"`
+	ToolCallID string     `json:"tool_call_id,omitempty"`
+	ToolCalls  []ToolCall `json:"tool_calls,omitempty"`
+	StopReason StopReason `json:"stop_reason,omitempty"`
+	IsError    bool       `json:"is_error,omitempty"`
+}
+
+// ToolDecl is the JSON-schema declaration sent to the provider.
+type ToolDecl struct {
+	Name        string          `json:"name"`
+	Description string          `json:"description"`
+	Parameters  json.RawMessage `json:"parameters"`
+}
+
+// Usage counts tokens reported by the provider. Zero means the provider omitted usage.
+type Usage struct {
+	PromptTokens     int `json:"prompt_tokens"`
+	CompletionTokens int `json:"completion_tokens"`
+	TotalTokens      int `json:"total_tokens"`
+}
+
+// Event is one streamed update. Failures are KindError plus a final assistant
+// message whose StopReason is error or aborted. Stream implementations do not panic.
+type Event struct {
+	Kind           EventKind
+	Text           string
+	ToolIndex      int
+	ToolID         string
+	ToolName       string
+	ArgumentsDelta string
+	Message        Message
+	Usage          Usage
+	Err            error
+	Retryable      bool
+	ErrorMessage   string
+}
+
+// StreamOptions tunes one provider request.
+type StreamOptions struct {
+	MaxTokens int
+	APIBase   string
+	APIKey    string
+}
+
+// StreamFunc is the provider boundary. It must return a channel that is closed
+// after a terminal message_end or error event. It must not panic.
+type StreamFunc func(ctx context.Context, model string, messages []Message, tools []ToolDecl, opts StreamOptions) <-chan Event