update
This commit is contained in:
1 parent
a7d8b6107a
commit
9f578780bd
2 files changed
+1144
No files matched your search
+773
@@ -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 = ¶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: <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)
|
||||||
|
}
|
||||||
@@ -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)
|
||||||
|
}
|
||||||
|
}
|
||||||
Reference in new issue
Block a user