diff --git a/mcp/agent.go b/mcp/agent.go new file mode 100644 index 0000000..bebbc9f --- /dev/null +++ b/mcp/agent.go @@ -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), (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 " 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 = ¶llel + 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: \ndata: \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: \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) +} diff --git a/mcp/agent_test.go b/mcp/agent_test.go new file mode 100644 index 0000000..dde482e --- /dev/null +++ b/mcp/agent_test.go @@ -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) + } +}