This commit is contained in:
mats committed 2026-08-22 20:51:37 +02:00
1 parent a7d8b6107a
commit 9f578780bd
2 files changed
+1144

No files matched your search

+773
View File
@@ -0,0 +1,773 @@
package mcp
// Agent is the generic OpenAI↔MCP bridge: an OpenAI-compatible chat
// completions endpoint whose tools are discovered from — and executed against —
// any MCP endpoint. It knows nothing about content, carts, or orders:
//
// 1. On every request it calls tools/list on the MCP endpoint and advertises
// those tools to the model (merged with any client-supplied tools).
// 2. It runs the agentic loop, streaming text deltas to the client as SSE
// while intercepting tool calls and dispatching them via tools/call,
// appending the configured Bearer token.
// 3. Tool calls the MCP server does not know are passed through to the client
// (frontend tool calls) so the browser can execute them itself.
//
// The wire contract for streaming clients: named SSE events "delta" (a JSON
// string chunk), <ProgressEvent> (tool activity), "response" (final payload),
// then "[DONE]". Non-streaming clients get a normal chat.completion JSON body.
import (
"bufio"
"bytes"
"context"
"encoding/json"
"fmt"
"io"
"log/slog"
"net/http"
"strings"
"time"
)
// AgentConfig configures an Agent.
type AgentConfig struct {
// UpstreamURL is the base URL of an OpenAI-compatible API, e.g.
// "https://api.openai.com" (no trailing slash; /v1/chat/completions is appended).
UpstreamURL string
// UpstreamAPIKey is the Bearer token sent to the upstream AI endpoint.
UpstreamAPIKey string
// Model is the default model used when the request omits one.
Model string
// MCPURL is the streamable-HTTP MCP endpoint providing the tools.
MCPURL string
// MCPToken is appended as "Authorization: Bearer <MCPToken>" on every MCP
// request. Empty means no Authorization header (e.g. a loopback listener).
MCPToken string
// SystemPrompt is prepended as a system message unless the conversation
// already starts with one. Empty disables injection.
SystemPrompt string
// ProgressEvent is the SSE event name for tool-activity events
// (default "tool_progress").
ProgressEvent string
// MaxToolIterations caps the agentic loop (default 10).
MaxToolIterations int
// LocalTools are executed in-process by the agent instead of being
// dispatched over the MCP endpoint. They are advertised to the model
// alongside the MCP tools but never reach MCP clients — for capabilities
// that are agent-only (e.g. a CDP browser tool returning base64
// screenshots too bulky and too privileged for the shared tool surface).
LocalTools []Tool
// Logger receives structured operational logs. Nil uses slog.Default().
Logger *slog.Logger
}
// Agent is the HTTP handler for the chat completions endpoint. Build it with
// NewAgent and mount Handler().
type Agent struct {
cfg AgentConfig
client *http.Client
log *slog.Logger
}
// NewAgent builds an Agent.
func NewAgent(cfg AgentConfig) *Agent {
if cfg.ProgressEvent == "" {
cfg.ProgressEvent = "tool_progress"
}
if cfg.MaxToolIterations <= 0 {
cfg.MaxToolIterations = 10
}
if cfg.Logger == nil {
cfg.Logger = slog.Default()
}
return &Agent{
cfg: cfg,
client: &http.Client{Timeout: 5 * time.Minute},
log: cfg.Logger,
}
}
// Enabled reports whether the agent has enough configuration to serve.
func (a *Agent) Enabled() bool { return a.cfg.UpstreamURL != "" && a.cfg.UpstreamAPIKey != "" && a.cfg.MCPURL != "" }
// Handler returns the http.Handler to mount at the chat completions route.
func (a *Agent) Handler() http.Handler { return http.HandlerFunc(a.serveHTTP) }
// ── OpenAI wire types ──────────────────────────────────────────────────────────
// Message is one turn in the conversation.
type Message struct {
Role string `json:"role"`
Content any `json:"content,omitempty"` // string, null, or content-part array
ToolCalls []ToolCall `json:"tool_calls,omitempty"`
ToolCallID string `json:"tool_call_id,omitempty"`
}
// ToolCall is a single function invocation requested by the model.
type ToolCall struct {
ID string `json:"id"`
Type string `json:"type"` // "function"
Function FunctionCall `json:"function"`
}
// FunctionCall holds the function name and its JSON-serialised arguments.
type FunctionCall struct {
Name string `json:"name"`
Arguments string `json:"arguments"`
}
// ChatRequest is the OpenAI-compatible request body.
type ChatRequest struct {
Model string `json:"model"`
Messages []Message `json:"messages"`
Tools []ToolDef `json:"tools,omitempty"`
ToolChoice any `json:"tool_choice,omitempty"`
Stream bool `json:"stream"`
Temperature *float64 `json:"temperature,omitempty"`
MaxTokens *int `json:"max_tokens,omitempty"`
TopP *float64 `json:"top_p,omitempty"`
ParallelToolCalls *bool `json:"parallel_tool_calls,omitempty"`
}
// ToolDef is an OpenAI function-tool definition.
type ToolDef struct {
Type string `json:"type"` // "function"
Function FuncDef `json:"function"`
}
// FuncDef is the function schema advertised to the model.
type FuncDef struct {
Name string `json:"name"`
Description string `json:"description"`
Parameters json.RawMessage `json:"parameters"`
}
// ChatResponse is the non-streaming response body.
type ChatResponse struct {
ID string `json:"id"`
Object string `json:"object"`
Created int64 `json:"created"`
Model string `json:"model"`
Choices []Choice `json:"choices"`
}
// Choice is one completion candidate.
type Choice struct {
Index int `json:"index"`
Message Message `json:"message"`
FinishReason string `json:"finish_reason"`
}
// FrontendToolCall is the shape passed through to browser-side executors:
// {id, name, arguments} instead of {id, type, function:{name, arguments}}.
type FrontendToolCall struct {
ID string `json:"id"`
Name string `json:"name"`
Arguments any `json:"arguments"`
}
// toolProgress is emitted as a named SSE event for every tool invocation so
// streaming clients can show a live activity log.
type toolProgress struct {
Iter int `json:"iter"`
Tool string `json:"tool"`
Args string `json:"args,omitempty"`
Status string `json:"status"` // "calling" | "done" | "error"
Err string `json:"error,omitempty"`
Bytes int `json:"bytes,omitempty"`
Result any `json:"result,omitempty"`
}
// ── HTTP handler ───────────────────────────────────────────────────────────────
func (a *Agent) serveHTTP(w http.ResponseWriter, r *http.Request) {
if r.Method != http.MethodPost {
w.Header().Set("Allow", "POST")
http.Error(w, "method not allowed", http.StatusMethodNotAllowed)
return
}
body, err := io.ReadAll(io.LimitReader(r.Body, 4<<20))
if err != nil {
http.Error(w, "read error", http.StatusBadRequest)
return
}
var req ChatRequest
if err := json.Unmarshal(body, &req); err != nil {
http.Error(w, "invalid request: "+err.Error(), http.StatusBadRequest)
return
}
if req.Model == "" {
req.Model = a.cfg.Model
}
mcpTools, err := a.listTools(r.Context())
if err != nil {
a.log.Error("agent: tools/list failed", "err", err)
http.Error(w, "mcp tools/list: "+err.Error(), http.StatusBadGateway)
return
}
req = mergeTools(req, mcpTools, localToolDefs(a.cfg.LocalTools))
if a.cfg.SystemPrompt != "" && (len(req.Messages) == 0 || req.Messages[0].Role != "system") {
req.Messages = append([]Message{{Role: "system", Content: a.cfg.SystemPrompt}}, req.Messages...)
}
a.log.Info("agent: request",
"model", req.Model,
"stream", req.Stream,
"messages", len(req.Messages),
"tools", len(req.Tools),
)
var flush func()
if req.Stream {
w.Header().Set("Content-Type", "text/event-stream")
w.Header().Set("Cache-Control", "no-cache")
w.Header().Set("Connection", "keep-alive")
w.Header().Set("X-Accel-Buffering", "no")
w.WriteHeader(http.StatusOK)
if f, ok := w.(http.Flusher); ok {
flush = f.Flush
} else {
flush = func() {}
}
}
sendErr := func(status int, msg string) {
if flush != nil {
emitSSEEvent(w, "error", map[string]string{"error": msg})
emitSSE(w, "[DONE]")
flush()
return
}
http.Error(w, msg, status)
}
content, finishReason, frontendCalls, ok := a.loop(r.Context(), w, flush, sendErr, &req)
if !ok {
return // sendErr already wrote the response
}
if !req.Stream {
msg := Message{Role: "assistant", Content: content}
if len(frontendCalls) > 0 {
for _, fc := range frontendCalls {
args, _ := json.Marshal(fc.Arguments)
msg.ToolCalls = append(msg.ToolCalls, ToolCall{
ID: fc.ID,
Type: "function",
Function: FunctionCall{Name: fc.Name, Arguments: string(args)},
})
}
}
writeJSON(w, http.StatusOK, ChatResponse{
ID: "chatcmpl-agent",
Object: "chat.completion",
Created: time.Now().Unix(),
Model: req.Model,
Choices: []Choice{{Index: 0, Message: msg, FinishReason: finishReason}},
})
return
}
emitSSEEvent(w, "response", map[string]any{
"content": content,
"tool_calls": frontendCalls,
"finish_reason": finishReason,
})
emitSSE(w, "[DONE]")
flush()
}
// loop runs the agentic loop until the model produces a final answer or issues
// a frontend tool call. On upstream failure it calls sendErr and returns ok=false.
func (a *Agent) loop(
ctx context.Context,
w io.Writer,
flush func(),
sendErr func(status int, msg string),
req *ChatRequest,
) (content, finishReason string, frontendCalls []FrontendToolCall, ok bool) {
parallel := true
req.ParallelToolCalls = &parallel
messages := req.Messages
onDelta := func(delta string) {
if flush != nil && delta != "" {
b, _ := json.Marshal(delta)
fmt.Fprintf(w, "event: delta\ndata: %s\n\n", b)
flush()
}
}
for iter := 0; iter < a.cfg.MaxToolIterations; iter++ {
reqCopy := *req
reqCopy.Messages = messages
reqCopy.Stream = true // always stream from upstream so deltas forward live
resp, err := a.callUpstream(ctx, reqCopy)
if err != nil {
a.log.Error("agent: upstream call failed", "iter", iter, "err", err)
sendErr(http.StatusBadGateway, "upstream error: "+err.Error())
return "", "", nil, false
}
if resp.StatusCode != http.StatusOK {
respBody, _ := io.ReadAll(io.LimitReader(resp.Body, 8<<20))
resp.Body.Close()
a.log.Error("agent: upstream non-200", "iter", iter, "status", resp.StatusCode)
sendErr(resp.StatusCode, "upstream error: "+string(respBody))
return "", "", nil, false
}
content, toolCalls, finishReason, err := readUpstreamStream(ctx, resp, onDelta)
if err != nil {
a.log.Error("agent: upstream stream error", "iter", iter, "err", err)
sendErr(http.StatusBadGateway, "upstream stream error: "+err.Error())
return "", "", nil, false
}
a.log.Info("agent: upstream done",
"iter", iter, "finish_reason", finishReason,
"tool_calls", len(toolCalls), "content_len", len(content),
)
msg := Message{Role: "assistant", Content: content, ToolCalls: toolCalls}
messages = append(messages, msg)
if finishReason != "tool_calls" || len(toolCalls) == 0 {
return content, finishReason, nil, true
}
// Execute known tools (local in-process first, then MCP); pass unknown
// ones back to the client.
local := localToolIndex(a.cfg.LocalTools)
var remaining []ToolCall
for _, tc := range toolCalls {
a.emitProgress(w, flush, toolProgress{
Iter: iter, Tool: tc.Function.Name, Args: tc.Function.Arguments, Status: "calling",
})
var (
resultContent string
isError bool
callErr error
)
if t, ok := local[tc.Function.Name]; ok {
// In-process execution — never touches the MCP endpoint.
out, invokeErr := t.Invoke(ctx, json.RawMessage(tc.Function.Arguments))
if invokeErr != nil {
callErr = invokeErr
resultContent = fmt.Sprintf(`{"error":%q}`, invokeErr.Error())
a.log.Warn("agent: local tool error", "tool", tc.Function.Name, "err", invokeErr)
} else {
b, _ := json.Marshal(out)
resultContent = string(b)
}
} else {
out, found, callToolErr := a.callTool(ctx, tc.Function.Name, tc.Function.Arguments)
if !found {
remaining = append(remaining, tc)
continue
}
resultContent, isError, callErr = out.text, out.isError, callToolErr
}
status, summary := "done", any(nil)
if callErr != nil {
status = "error"
a.log.Warn("agent: tool error", "tool", tc.Function.Name, "err", callErr)
} else if isError {
status = "error"
} else {
summary = previewOf(resultContent)
}
a.emitProgress(w, flush, toolProgress{
Iter: iter, Tool: tc.Function.Name, Status: status,
Bytes: len(resultContent), Result: summary, Err: errText(callErr),
})
messages = append(messages, Message{
Role: "tool", Content: resultContent, ToolCallID: tc.ID,
})
}
if len(remaining) > 0 {
ftc := make([]FrontendToolCall, len(remaining))
for i, tc := range remaining {
var args any
_ = json.Unmarshal([]byte(tc.Function.Arguments), &args)
ftc[i] = FrontendToolCall{ID: tc.ID, Name: tc.Function.Name, Arguments: args}
}
return content, finishReason, ftc, true
}
}
a.log.Error("agent: agentic loop did not converge", "maxIter", a.cfg.MaxToolIterations)
sendErr(http.StatusInternalServerError, "agentic loop did not converge")
return "", "", nil, false
}
func (a *Agent) emitProgress(w io.Writer, flush func(), p toolProgress) {
if flush == nil {
return
}
emitSSEEvent(w, a.cfg.ProgressEvent, p)
flush()
}
func errText(err error) string {
if err == nil {
return ""
}
return err.Error()
}
// mergeTools injects the agent-side tools into the request, deduplicating
// against tools the client already supplied.
func mergeTools(req ChatRequest, toolSets ...[]ToolDef) ChatRequest {
existing := make(map[string]bool, len(req.Tools))
for _, t := range req.Tools {
existing[t.Function.Name] = true
}
for _, set := range toolSets {
for _, t := range set {
if !existing[t.Function.Name] {
req.Tools = append(req.Tools, t)
}
}
}
return req
}
// localToolDefs converts in-process tools to OpenAI definitions.
func localToolDefs(tools []Tool) []ToolDef {
defs := make([]ToolDef, 0, len(tools))
for _, t := range tools {
schema := t.InputSchema
if len(schema) == 0 {
schema = Object(nil, nil)
}
defs = append(defs, ToolDef{Type: "function", Function: FuncDef{
Name: t.Name,
Description: t.Description,
Parameters: schema,
}})
}
return defs
}
// ── Upstream (OpenAI-compatible) client ────────────────────────────────────────
func (a *Agent) callUpstream(ctx context.Context, req ChatRequest) (*http.Response, error) {
b, err := json.Marshal(req)
if err != nil {
return nil, fmt.Errorf("marshal request: %w", err)
}
httpReq, err := http.NewRequestWithContext(ctx, http.MethodPost,
a.cfg.UpstreamURL+"/v1/chat/completions", bytes.NewReader(b))
if err != nil {
return nil, fmt.Errorf("build request: %w", err)
}
httpReq.Header.Set("Content-Type", "application/json")
httpReq.Header.Set("Authorization", "Bearer "+a.cfg.UpstreamAPIKey)
return a.client.Do(httpReq)
}
// readUpstreamStream reads an upstream SSE chat stream, forwarding each content
// delta to onDelta, and returns the accumulated content, completed tool calls,
// and finish reason.
func readUpstreamStream(ctx context.Context, resp *http.Response, onDelta func(string)) (content string, toolCalls []ToolCall, finishReason string, err error) {
defer resp.Body.Close()
reader := bufio.NewReader(resp.Body)
type partialTool struct {
id string
name string
arguments strings.Builder
}
partials := make(map[int]*partialTool)
var contentBuf strings.Builder
for {
select {
case <-ctx.Done():
return contentBuf.String(), nil, "", ctx.Err()
default:
}
line, readErr := reader.ReadString('\n')
if readErr != nil {
if readErr == io.EOF {
break
}
return contentBuf.String(), nil, "", readErr
}
line = strings.TrimSpace(line)
if !strings.HasPrefix(line, "data: ") {
continue
}
data := strings.TrimPrefix(line, "data: ")
if data == "[DONE]" {
break
}
var chunk struct {
Choices []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"`
} `json:"choices"`
}
if err := json.Unmarshal([]byte(data), &chunk); err != nil {
continue
}
if len(chunk.Choices) == 0 {
continue
}
choice := chunk.Choices[0]
if choice.Delta.Content != "" {
contentBuf.WriteString(choice.Delta.Content)
onDelta(choice.Delta.Content)
}
for _, tc := range choice.Delta.ToolCalls {
p, ok := partials[tc.Index]
if !ok {
p = &partialTool{}
partials[tc.Index] = p
}
if tc.ID != "" {
p.id = tc.ID
}
if tc.Function.Name != "" {
p.name = tc.Function.Name
}
if tc.Function.Arguments != "" {
p.arguments.WriteString(tc.Function.Arguments)
}
}
if choice.FinishReason != nil {
finishReason = *choice.FinishReason
break
}
}
if finishReason == "tool_calls" {
for i := 0; ; i++ {
p, ok := partials[i]
if !ok {
break
}
toolCalls = append(toolCalls, ToolCall{
ID: p.id,
Type: "function",
Function: FunctionCall{Name: p.name, Arguments: p.arguments.String()},
})
}
}
return contentBuf.String(), toolCalls, finishReason, nil
}
// localToolIndex builds a name→tool map for O(1) in-process dispatch.
func localToolIndex(tools []Tool) map[string]Tool {
if len(tools) == 0 {
return nil
}
idx := make(map[string]Tool, len(tools))
for _, t := range tools {
idx[t.Name] = t
}
return idx
}
// ── MCP client (tools/list + tools/call over streamable HTTP) ─────────────────
type mcpToolList struct {
Tools []struct {
Name string `json:"name"`
Description string `json:"description"`
InputSchema json.RawMessage `json:"inputSchema"`
} `json:"tools"`
}
type mcpCallResult struct {
Content []struct {
Type string `json:"type"`
Text string `json:"text"`
} `json:"content"`
IsError bool `json:"isError"`
}
// toolOutput is the normalised result of a tools/call.
type toolOutput struct {
text string // joined text blocks — what the model sees
isError bool // MCP isError flag (a tool-level failure the model can read)
}
// listTools fetches the MCP server's advertised tools as OpenAI definitions.
func (a *Agent) listTools(ctx context.Context) ([]ToolDef, error) {
var out mcpToolList
if err := a.mcpRequest(ctx, "tools/list", nil, &out); err != nil {
return nil, err
}
defs := make([]ToolDef, 0, len(out.Tools))
for _, t := range out.Tools {
schema := t.InputSchema
if len(schema) == 0 {
schema = Object(nil, nil)
}
defs = append(defs, ToolDef{Type: "function", Function: FuncDef{
Name: t.Name,
Description: t.Description,
Parameters: schema,
}})
}
return defs, nil
}
// callTool dispatches one tools/call. found=false means the MCP server does not
// know the tool (the caller passes it through to the frontend).
func (a *Agent) callTool(ctx context.Context, name, argsJSON string) (out toolOutput, found bool, err error) {
params := map[string]any{"name": name}
if strings.TrimSpace(argsJSON) != "" {
var args any
if json.Unmarshal([]byte(argsJSON), &args) == nil {
params["arguments"] = args
}
}
var res mcpCallResult
if err := a.mcpRequest(ctx, "tools/call", params, &res); err != nil {
if strings.Contains(err.Error(), "unknown tool") {
return toolOutput{}, false, nil
}
return toolOutput{}, true, err
}
var texts strings.Builder
for _, c := range res.Content {
if c.Text != "" {
if texts.Len() > 0 {
texts.WriteString("\n")
}
texts.WriteString(c.Text)
}
}
return toolOutput{
text: texts.String(),
isError: res.IsError,
}, true, nil
}
// previewOf builds a tiny structured preview of a tool's JSON output for the
// progress event: top-level keys for objects, length for arrays.
func previewOf(s string) any {
var raw any
if json.Unmarshal([]byte(s), &raw) != nil {
return nil
}
switch v := raw.(type) {
case map[string]any:
keys := make([]string, 0, len(v))
for k := range v {
keys = append(keys, k)
}
return map[string]any{"keys": keys}
case []any:
return map[string]any{"count": len(v)}
default:
return nil
}
}
// mcpRequest posts one JSON-RPC request to the MCP endpoint, appending the
// configured Bearer token, and decodes the result into out.
func (a *Agent) mcpRequest(ctx context.Context, method string, params, out any) error {
body, err := json.Marshal(map[string]any{
"jsonrpc": "2.0",
"id": 1,
"method": method,
"params": params,
})
if err != nil {
return err
}
req, err := http.NewRequestWithContext(ctx, http.MethodPost, a.cfg.MCPURL, bytes.NewReader(body))
if err != nil {
return err
}
req.Header.Set("Content-Type", "application/json")
req.Header.Set("Accept", "application/json")
if a.cfg.MCPToken != "" {
req.Header.Set("Authorization", "Bearer "+a.cfg.MCPToken)
}
resp, err := http.DefaultClient.Do(req)
if err != nil {
return fmt.Errorf("mcp %s: %w", method, err)
}
defer resp.Body.Close()
respBody, err := io.ReadAll(io.LimitReader(resp.Body, 16<<20))
if err != nil {
return fmt.Errorf("mcp %s: read: %w", method, err)
}
if resp.StatusCode != http.StatusOK {
return fmt.Errorf("mcp %s: status %d: %s", method, resp.StatusCode, truncate(string(respBody), 512))
}
var rpcResp struct {
Result json.RawMessage `json:"result"`
Error *struct {
Code int `json:"code"`
Message string `json:"message"`
} `json:"error"`
}
if err := json.Unmarshal(respBody, &rpcResp); err != nil {
return fmt.Errorf("mcp %s: decode response: %w", method, err)
}
if rpcResp.Error != nil {
return fmt.Errorf("mcp %s: %s", method, rpcResp.Error.Message)
}
if out != nil && len(rpcResp.Result) > 0 {
if err := json.Unmarshal(rpcResp.Result, out); err != nil {
return fmt.Errorf("mcp %s: decode result: %w", method, err)
}
}
return nil
}
func truncate(s string, n int) string {
if len(s) <= n {
return s
}
return s[:n] + "…"
}
// ── SSE helpers ────────────────────────────────────────────────────────────────
// emitSSEEvent writes a named SSE event: "event: <name>\ndata: <json>\n\n".
func emitSSEEvent(w io.Writer, event string, v any) {
fmt.Fprintf(w, "event: %s\n", event)
emitSSE(w, v)
}
// emitSSE writes one SSE data line: "data: <payload>\n\n".
func emitSSE(w io.Writer, v any) {
payload, ok := v.(string)
if !ok {
b, _ := json.Marshal(v)
payload = string(b)
}
fmt.Fprintf(w, "data: %s\n\n", payload)
}
// writeJSON encodes v as JSON with the given status.
func writeJSON(w http.ResponseWriter, status int, v any) {
w.Header().Set("Content-Type", "application/json")
w.WriteHeader(status)
_ = json.NewEncoder(w).Encode(v)
}
+371
View File
@@ -0,0 +1,371 @@
package mcp
import (
"context"
"encoding/json"
"fmt"
"net/http"
"net/http/httptest"
"strings"
"testing"
)
// stubMCP serves a minimal streamable-HTTP MCP server with one "echo" tool and
// records the Authorization header of the last request.
func stubMCP(t *testing.T, lastAuth *string) *httptest.Server {
t.Helper()
return httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
*lastAuth = r.Header.Get("Authorization")
var req struct {
Method string `json:"method"`
Params json.RawMessage `json:"params"`
}
if err := json.NewDecoder(r.Body).Decode(&req); err != nil {
t.Errorf("decode rpc: %v", err)
return
}
var result any
switch req.Method {
case "tools/list":
result = map[string]any{"tools": []map[string]any{{
"name": "echo",
"description": "Echo back its input text",
"inputSchema": Object(Props{"text": String("the text to echo")}, []string{"text"}),
}}}
case "tools/call":
var p struct {
Name string `json:"name"`
Args struct {
Text string `json:"text"`
} `json:"arguments"`
}
_ = json.Unmarshal(req.Params, &p)
if p.Name != "echo" {
http.Error(w, "unknown tool", http.StatusBadRequest)
return
}
result = map[string]any{
"content": []map[string]any{{"type": "text", "text": `{"echo":"` + p.Args.Text + `"}`}},
}
default:
http.Error(w, "unknown method", http.StatusBadRequest)
return
}
b, _ := json.Marshal(map[string]any{"jsonrpc": "2.0", "id": 1, "result": result})
w.Header().Set("Content-Type", "application/json")
_, _ = w.Write(b)
}))
}
// stubUpstream serves an OpenAI-compatible /v1/chat/completions SSE stream.
// On the first call it emits one echo tool call; afterwards it streams
// "done " + the tool message content as the final answer.
func stubUpstream(t *testing.T, calls *int, sawTools *[][]string) *httptest.Server {
t.Helper()
mux := http.NewServeMux()
mux.HandleFunc("POST /v1/chat/completions", func(w http.ResponseWriter, r *http.Request) {
var req ChatRequest
if err := json.NewDecoder(r.Body).Decode(&req); err != nil {
t.Errorf("decode chat request: %v", err)
return
}
if sawTools != nil {
names := make([]string, 0, len(req.Tools))
for _, tl := range req.Tools {
names = append(names, tl.Function.Name)
}
*sawTools = append(*sawTools, names)
}
*calls++
w.Header().Set("Content-Type", "text/event-stream")
flusher := w.(http.Flusher)
send := func(v any) {
b, _ := json.Marshal(v)
fmt.Fprintf(w, "data: %s\n\n", b)
flusher.Flush()
}
if *calls == 1 {
send(struct {
Choices []struct {
Delta struct {
ToolCalls []ToolCall `json:"tool_calls"`
} `json:"delta"`
FinishReason *string `json:"finish_reason"`
} `json:"choices"`
}{Choices: []struct {
Delta struct {
ToolCalls []ToolCall `json:"tool_calls"`
} `json:"delta"`
FinishReason *string `json:"finish_reason"`
}{{
Delta: struct {
ToolCalls []ToolCall `json:"tool_calls"`
}{ToolCalls: []ToolCall{{
ID: "call_1",
Type: "function",
Function: FunctionCall{Name: "echo", Arguments: `{"text":"hi"}`},
}}},
FinishReason: strPtr("tool_calls"),
}}})
} else {
// Second round: the model should have seen the tool result message.
toolMsg := ""
for _, m := range req.Messages {
if m.Role == "tool" && m.ToolCallID == "call_1" {
toolMsg, _ = m.Content.(string)
}
}
send(struct {
Choices []struct {
Delta struct {
Content string `json:"content"`
} `json:"delta"`
FinishReason *string `json:"finish_reason"`
} `json:"choices"`
}{Choices: []struct {
Delta struct {
Content string `json:"content"`
} `json:"delta"`
FinishReason *string `json:"finish_reason"`
}{{
Delta: struct {
Content string `json:"content"`
}{Content: "saw:" + toolMsg},
FinishReason: strPtr("stop"),
}}})
}
fmt.Fprint(w, "data: [DONE]\n\n")
flusher.Flush()
})
return httptest.NewServer(mux)
}
func strPtr(s string) *string { return &s }
func TestAgentAgenticLoop(t *testing.T) {
var lastAuth string
mcpSrv := stubMCP(t, &lastAuth)
defer mcpSrv.Close()
var calls int
var sawTools [][]string
upSrv := stubUpstream(t, &calls, &sawTools)
defer upSrv.Close()
agent := NewAgent(AgentConfig{
UpstreamURL: upSrv.URL,
UpstreamAPIKey: "sk-test",
MCPURL: mcpSrv.URL + "/mcp",
MCPToken: "mcp-secret",
SystemPrompt: "You are a test.",
ProgressEvent: "cms_progress",
})
body := `{"model":"test-model","messages":[{"role":"user","content":"hi"}],"stream":true}`
req := httptest.NewRequest(http.MethodPost, "/chat/completions", strings.NewReader(body))
rec := httptest.NewRecorder()
agent.Handler().ServeHTTP(rec, req)
if rec.Code != http.StatusOK {
t.Fatalf("status = %d, body = %s", rec.Code, rec.Body.String())
}
sse := rec.Body.String()
// The model must have been offered the MCP tool (round 1)…
if len(sawTools) < 2 || sawTools[0][0] != "echo" {
t.Fatalf("model was not offered the echo tool: %v", sawTools)
}
// …and executed via tools/call with the bearer token appended.
if lastAuth != "Bearer mcp-secret" {
t.Errorf("MCP Authorization = %q, want Bearer mcp-secret", lastAuth)
}
// The final answer must reflect the tool result fed back to the model.
if !strings.Contains(sse, `"saw:{\"echo\":\"hi\"}"`) {
t.Errorf("final delta missing tool feedback; SSE:\n%s", sse)
}
// Wire contract: progress event, response event, [DONE].
for _, want := range []string{"event: cms_progress", `"status":"calling"`, `"status":"done"`, "event: response", `"finish_reason":"stop"`, "data: [DONE]"} {
if !strings.Contains(sse, want) {
t.Errorf("SSE missing %q;\n%s", want, sse)
}
}
if calls != 2 {
t.Errorf("upstream calls = %d, want 2", calls)
}
}
func TestAgentLocalTool(t *testing.T) {
var lastAuth string
mcpSrv := stubMCP(t, &lastAuth)
defer mcpSrv.Close()
// Upstream asks for the local tool; second round confirms the model saw
// its in-process result.
var calls int
upSrv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
calls++
w.Header().Set("Content-Type", "text/event-stream")
flusher := w.(http.Flusher)
send := func(v any) {
b, _ := json.Marshal(v)
fmt.Fprintf(w, "data: %s\n\n", b)
flusher.Flush()
}
if calls == 1 {
send(struct {
Choices []struct {
Delta struct {
ToolCalls []ToolCall `json:"tool_calls"`
} `json:"delta"`
FinishReason *string `json:"finish_reason"`
} `json:"choices"`
}{Choices: []struct {
Delta struct {
ToolCalls []ToolCall `json:"tool_calls"`
} `json:"delta"`
FinishReason *string `json:"finish_reason"`
}{{
Delta: struct {
ToolCalls []ToolCall `json:"tool_calls"`
}{ToolCalls: []ToolCall{{
ID: "loc_1",
Type: "function",
Function: FunctionCall{Name: "local_ping", Arguments: `{}`},
}}},
FinishReason: strPtr("tool_calls"),
}}})
} else {
toolMsg := ""
for _, m := range reqMessages(r) {
if m.Role == "tool" && m.ToolCallID == "loc_1" {
toolMsg, _ = m.Content.(string)
}
}
send(struct {
Choices []struct {
Delta struct {
Content string `json:"content"`
} `json:"delta"`
FinishReason *string `json:"finish_reason"`
} `json:"choices"`
}{Choices: []struct {
Delta struct {
Content string `json:"content"`
} `json:"delta"`
FinishReason *string `json:"finish_reason"`
}{{
Delta: struct {
Content string `json:"content"`
}{Content: toolMsg},
FinishReason: strPtr("stop"),
}}})
}
fmt.Fprint(w, "data: [DONE]\n\n")
flusher.Flush()
}))
defer upSrv.Close()
authBefore := lastAuth
agent := NewAgent(AgentConfig{
UpstreamURL: upSrv.URL, UpstreamAPIKey: "k",
MCPURL: mcpSrv.URL + "/mcp",
LocalTools: []Tool{{
Name: "local_ping",
Description: "In-process ping",
InputSchema: Object(nil, nil),
Invoke: func(ctx context.Context, args json.RawMessage) (any, error) {
return map[string]any{"pong": true}, nil
},
}},
})
req := httptest.NewRequest(http.MethodPost, "/chat/completions",
strings.NewReader(`{"messages":[{"role":"user","content":"ping"}],"stream":true}`))
rec := httptest.NewRecorder()
agent.Handler().ServeHTTP(rec, req)
sse := rec.Body.String()
if !strings.Contains(sse, `\"pong\":true`) {
t.Errorf("local tool result not fed back to the model:\n%s", sse)
}
if lastAuth != authBefore {
t.Errorf("local tool dispatch leaked to the MCP endpoint (auth header changed)")
}
}
// reqMessages decodes the chat request body from a recorded upstream call.
func reqMessages(r *http.Request) []Message {
var req ChatRequest
_ = json.NewDecoder(r.Body).Decode(&req)
return req.Messages
}
func TestAgentFrontendToolCallPassthrough(t *testing.T) {
var lastAuth string
mcpSrv := stubMCP(t, &lastAuth)
defer mcpSrv.Close()
calls := 0
// Upstream immediately requests an unknown (frontend) tool.
upSrv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
calls++
w.Header().Set("Content-Type", "text/event-stream")
fmt.Fprintf(w, `data: {"choices":[{"delta":{"tool_calls":[{"index":0,"id":"fe_1","type":"function","function":{"name":"browser_exec","arguments":"{\"x\":1}"}}]},"finish_reason":"tool_calls"}]}`+"\n\n")
fmt.Fprint(w, "data: [DONE]\n\n")
}))
defer upSrv.Close()
agent := NewAgent(AgentConfig{
UpstreamURL: upSrv.URL, UpstreamAPIKey: "k",
MCPURL: mcpSrv.URL + "/mcp",
})
req := httptest.NewRequest(http.MethodPost, "/chat/completions",
strings.NewReader(`{"messages":[{"role":"user","content":"go"}],"stream":true}`))
rec := httptest.NewRecorder()
agent.Handler().ServeHTTP(rec, req)
if !strings.Contains(rec.Body.String(), `"name":"browser_exec"`) ||
!strings.Contains(rec.Body.String(), `"id":"fe_1"`) {
t.Errorf("frontend tool call not passed through:\n%s", rec.Body.String())
}
if calls != 1 {
t.Errorf("upstream calls = %d, want 1 (loop must stop for frontend calls)", calls)
}
}
func TestAgentNonStreaming(t *testing.T) {
var lastAuth string
mcpSrv := stubMCP(t, &lastAuth)
defer mcpSrv.Close()
var calls int
upSrv := stubUpstream(t, &calls, nil)
defer upSrv.Close()
agent := NewAgent(AgentConfig{
UpstreamURL: upSrv.URL, UpstreamAPIKey: "k",
MCPURL: mcpSrv.URL + "/mcp",
Model: "default-model",
})
req := httptest.NewRequest(http.MethodPost, "/chat/completions",
strings.NewReader(`{"messages":[{"role":"user","content":"hi"}]}`))
rec := httptest.NewRecorder()
agent.Handler().ServeHTTP(rec, req)
var resp ChatResponse
if err := json.Unmarshal(rec.Body.Bytes(), &resp); err != nil {
t.Fatalf("decode response: %v\n%s", err, rec.Body.String())
}
if len(resp.Choices) != 1 || resp.Choices[0].FinishReason != "stop" {
t.Fatalf("unexpected response: %+v", resp)
}
if !strings.Contains(resp.Choices[0].Message.Content.(string), `saw:{"echo":"hi"}`) {
t.Errorf("content = %q", resp.Choices[0].Message.Content)
}
if resp.Model != "default-model" {
t.Errorf("model default not applied: %q", resp.Model)
}
}