From 2336c9433cee5a6a068436b34f4cdb084dbfaba0 Mon Sep 17 00:00:00 2001 From: logan Date: Thu, 15 Oct 2026 13:00:40 +0800 Subject: [PATCH] =?UTF-8?q?feat(ai):=20Agent=20SSE=20+=20Anthropic/respons?= =?UTF-8?q?es=20tools=20+=20P1=20=E7=A1=AE=E8=AE=A4=E5=86=99?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit - Agent 支持 SSE（步骤边推、正文分块），JSON 非流式 fallback - tool calling 覆盖 openai-chat / anthropic / responses（单测解析回灌） - P1：mcp.json mode=read_write 才可写；需 confirm_danger；Agent 剥离模型自带令牌 - GET /api/ai/mcp-status、POST /api/ai/agent/confirm；审计 source=mcp|agent --- server/cmd/opsxdb-mcp/main.go | 1 + server/internal/ai/agent.go | 174 ++++++++++++-- server/internal/ai/agent_stream.go | 79 ++++++ server/internal/ai/agent_test.go | 2 +- server/internal/ai/ai_test.go | 100 ++++++++ server/internal/ai/config.go | 21 +- server/internal/ai/openai.go | 117 --------- server/internal/ai/remote.go | 292 ++++++++++++++++++++++ server/internal/ai/remote_tools.go | 307 +++++++++++++++++++++--- server/internal/ai/remote_tools_test.go | 72 ++++++ server/internal/aitools/audit.go | 1 + server/internal/aitools/policy.go | 6 +- server/internal/aitools/runtime.go | 85 ++++++- server/internal/aitools/runtime_test.go | 72 ++++++ server/internal/aitools/schema.go | 20 +- server/internal/aitools/status.go | 98 ++++++++ server/internal/server/router.go | 74 ++++++ 17 files changed, 1324 insertions(+), 197 deletions(-) create mode 100644 server/internal/ai/agent_stream.go create mode 100644 server/internal/ai/ai_test.go delete mode 100644 server/internal/ai/openai.go create mode 100644 server/internal/ai/remote.go create mode 100644 server/internal/ai/remote_tools_test.go create mode 100644 server/internal/aitools/status.go diff --git a/server/cmd/opsxdb-mcp/main.go b/server/cmd/opsxdb-mcp/main.go index d94d820..b3bbb65 100644 --- a/server/cmd/opsxdb-mcp/main.go +++ b/server/cmd/opsxdb-mcp/main.go @@ -26,6 +26,7 @@ func main() { fmt.Fprintf(os.Stderr, "opsxdb-mcp: load policy: %v\n", err) os.Exit(1) } + rt.Source = "mcp" s := server.NewMCPServer( "opsxdb", diff --git a/server/internal/ai/agent.go b/server/internal/ai/agent.go index b3f5fdf..da9ce1e 100644 --- a/server/internal/ai/agent.go +++ b/server/internal/ai/agent.go @@ -11,9 +11,10 @@ import ( ) const ( - defaultAgentMaxRounds = 8 - maxToolResultFeedChars = 12000 - maxStepArgsChars = 200 + defaultAgentMaxRounds = 8 + maxToolResultFeedChars = 12000 + maxStepArgsChars = 200 + maxGridRows = 200 ) // AgentReq 侧栏 Agent 请求（与 Ask 的 /api/ai/chat 分路径）。 @@ -24,16 +25,21 @@ type AgentReq struct { Database string `json:"database,omitempty"` Model string `json:"model,omitempty"` MaxRounds int `json:"maxRounds,omitempty"` + Stream bool `json:"stream,omitempty"` } // AgentStep 挂在回复下的步骤时间线一项。 type AgentStep struct { - Tool string `json:"tool"` - Args string `json:"args,omitempty"` - OK bool `json:"ok"` - Summary string `json:"summary"` - Error string `json:"error,omitempty"` - Ms int64 `json:"ms"` + Tool string `json:"tool"` + Args string `json:"args,omitempty"` + OK bool `json:"ok"` + Summary string `json:"summary"` + Error string `json:"error,omitempty"` + Ms int64 `json:"ms"` + NeedConfirm bool `json:"needConfirm,omitempty"` + ConfirmSQL string `json:"confirmSql,omitempty"` + Pending bool `json:"pending,omitempty"` + Result *AgentQueryResult `json:"result,omitempty"` } // AgentRes Agent 最终回答 + 步骤。 @@ -44,9 +50,18 @@ type AgentRes struct { Rounds int `json:"rounds"` } -// RunAgent 多轮：LLM tool_call → aitools.Runtime.Call → 回灌，直到最终文本或轮次耗尽。 -// ConfirmDanger 永不自动 true（由 Runtime 保证）。Ask 模式不要走这里。 +// RunAgent 多轮：LLM tool_call → aitools.Runtime.Call → 回灌。 +// ConfirmDanger 永不由模型自动带上（execOneTool 会剥离 confirm_danger）。 func RunAgent(ctx context.Context, llm ToolChat, rt *aitools.Runtime, req AgentReq) (AgentRes, error) { + return runAgent(ctx, llm, rt, req, nil) +} + +// RunAgentStream 同 RunAgent，经 emit 边跑边推（步骤 / 正文块 / done）。 +func RunAgentStream(ctx context.Context, llm ToolChat, rt *aitools.Runtime, req AgentReq, emit AgentEmitter) (AgentRes, error) { + return runAgent(ctx, llm, rt, req, emit) +} + +func runAgent(ctx context.Context, llm ToolChat, rt *aitools.Runtime, req AgentReq, emit AgentEmitter) (AgentRes, error) { if rt == nil { return AgentRes{}, fmt.Errorf("aitools runtime required") } @@ -65,6 +80,9 @@ func RunAgent(ctx context.Context, llm ToolChat, rt *aitools.Runtime, req AgentR maxRounds = 16 } + // Agent 路径禁止模型自行带确认令牌 + rt.RejectCallerConfirm = true + tools := aitools.OpenAITools() msgs := []Message{ {Role: "system", Content: buildAgentSystem(req)}, @@ -74,9 +92,15 @@ func RunAgent(ctx context.Context, llm ToolChat, rt *aitools.Runtime, req AgentR var steps []AgentStep rounds := 0 for round := 0; round < maxRounds; round++ { + if err := ctx.Err(); err != nil { + res := AgentRes{Steps: steps, Rounds: rounds, Content: "已取消"} + emitEvent(emit, AgentEvent{Type: AgentEventError, Error: "cancelled", Steps: steps, Rounds: rounds}) + return res, err + } rounds = round + 1 turn, err := llm.ChatWithTools(ctx, req.Model, msgs, tools) if err != nil { + emitEvent(emit, AgentEvent{Type: AgentEventError, Error: err.Error(), Steps: steps, Rounds: rounds}) return AgentRes{Steps: steps, Rounds: rounds}, err } if len(turn.ToolCalls) == 0 { @@ -84,15 +108,19 @@ func RunAgent(ctx context.Context, llm ToolChat, rt *aitools.Runtime, req AgentR if content == "" { content = "（模型未返回文本，也未调用工具）" } - return AgentRes{ + for _, part := range chunkDelta(content, 36) { + emitEvent(emit, AgentEvent{Type: AgentEventDelta, Delta: part}) + } + res := AgentRes{ Content: content, SQL: extractSQL(content), Steps: steps, Rounds: rounds, - }, nil + } + emitEvent(emit, AgentEvent{Type: AgentEventDone, Content: res.Content, SQL: res.SQL, Steps: steps, Rounds: rounds}) + return res, nil } - // assistant 消息带 tool_calls msgs = append(msgs, Message{ Role: "assistant", Content: turn.Content, @@ -100,28 +128,41 @@ func RunAgent(ctx context.Context, llm ToolChat, rt *aitools.Runtime, req AgentR }) for _, tc := range turn.ToolCalls { + name := strings.TrimSpace(tc.Function.Name) + pending := AgentStep{Tool: name, Args: truncateRunes(strings.TrimSpace(tc.Function.Arguments), maxStepArgsChars), Pending: true} + emitEvent(emit, AgentEvent{Type: AgentEventStepStart, Step: &pending}) + step, toolMsg := execOneTool(ctx, rt, tc) steps = append(steps, step) + cp := step + ev := AgentEvent{Type: AgentEventStep, Step: &cp} + if step.Result != nil { + ev.Result = step.Result + } + emitEvent(emit, ev) msgs = append(msgs, toolMsg) } } - return AgentRes{ - Content: fmt.Sprintf("已达最大工具轮次（%d）。请根据已执行步骤自行判断，或缩小问题后重试。", maxRounds), - Steps: steps, - Rounds: rounds, - }, nil + content := fmt.Sprintf("已达最大工具轮次（%d）。请根据已执行步骤自行判断，或缩小问题后重试。", maxRounds) + for _, part := range chunkDelta(content, 36) { + emitEvent(emit, AgentEvent{Type: AgentEventDelta, Delta: part}) + } + res := AgentRes{Content: content, Steps: steps, Rounds: rounds} + emitEvent(emit, AgentEvent{Type: AgentEventDone, Content: content, Steps: steps, Rounds: rounds}) + return res, nil } func buildAgentSystem(req AgentReq) string { var b strings.Builder - b.WriteString("你是 OpsxDB 数据库 Agent。你可以通过工具只读查询用户已白名单的连接。\n") + b.WriteString("你是 OpsxDB 数据库 Agent。你可以通过工具查询用户已白名单的连接。\n") b.WriteString("规则：\n") b.WriteString("1. 需要查库时必须调用工具；不要假装已执行 SQL。\n") - b.WriteString("2. 仅只读：execute_query 只允许 SELECT/SHOW/EXPLAIN/DESCRIBE/WITH…SELECT；Redis 仅读命令。\n") - b.WriteString("3. 先确认 connection_id（可用 opsxdb_list_connections）；用户给出偏好连接时优先用它。\n") - b.WriteString("4. 用中文简明回答；可执行 SQL 放在 ```sql 代码块中。\n") - b.WriteString("5. 工具结果可能截断，注意说明。\n") + b.WriteString("2. 默认只读：execute_query 以 SELECT/SHOW/EXPLAIN/DESCRIBE/WITH…SELECT 为主；Redis 以读命令为主。\n") + b.WriteString("3. 写操作若被拒绝并提示需要确认，向用户说明须在界面确认后执行；不要声称已写入。\n") + b.WriteString("4. 先确认 connection_id（可用 opsxdb_list_connections）；用户给出偏好连接时优先用它。\n") + b.WriteString("5. 用中文简明回答；可执行 SQL 放在 ```sql 代码块中。\n") + b.WriteString("6. 工具结果可能截断，注意说明。\n") if req.ConnectionID != "" { b.WriteString(fmt.Sprintf("\n用户当前偏好连接 id: %s", req.ConnectionID)) if req.Driver != "" { @@ -160,6 +201,9 @@ func execOneTool(ctx context.Context, rt *aitools.Runtime, tc ToolCall) (AgentSt Content: step.Error, } } + // 永不让模型自行带确认令牌 + delete(args, "confirm_danger") + delete(args, "confirmDanger") out, err := rt.Call(ctx, name, args) step.Ms = time.Since(start).Milliseconds() @@ -174,6 +218,37 @@ func execOneTool(ctx context.Context, rt *aitools.Runtime, tc ToolCall) (AgentSt Content: err.Error(), } } + + if m, ok := out.(map[string]any); ok { + if nc, _ := m["needConfirm"].(bool); nc { + step.OK = false + step.NeedConfirm = true + if sql, _ := m["sql"].(string); sql != "" { + step.ConfirmSQL = sql + } else if cmd, _ := m["command"].(string); cmd != "" { + step.ConfirmSQL = cmd + } + msg, _ := m["message"].(string) + if msg == "" { + msg = "requires user confirmation (confirm_danger)" + } + step.Error = msg + step.Summary = msg + feed := aitools.ResultJSON(out) + return step, Message{ + Role: "tool", + ToolCallID: tc.ID, + Name: name, + Content: feed, + } + } + if name == aitools.ToolExecuteQuery || name == aitools.ToolExecuteRedisCommand { + if qr := extractQueryResult(name, args, m); qr != nil { + step.Result = qr + } + } + } + step.OK = true step.Summary = summarizeToolResult(name, out) feed := aitools.ResultJSON(out) @@ -188,6 +263,57 @@ func execOneTool(ctx context.Context, rt *aitools.Runtime, tc ToolCall) (AgentSt } } +func extractQueryResult(tool string, args map[string]any, m map[string]any) *AgentQueryResult { + cols, _ := m["columns"].([]string) + if cols == nil { + if raw, ok := m["columns"].([]any); ok { + cols = make([]string, 0, len(raw)) + for _, c := range raw { + cols = append(cols, fmt.Sprint(c)) + } + } + } + var rows []map[string]interface{} + switch v := m["rows"].(type) { + case []map[string]interface{}: + rows = v + case []any: + for _, item := range v { + if row, ok := item.(map[string]interface{}); ok { + rows = append(rows, row) + } else if row, ok := item.(map[string]any); ok { + rows = append(rows, row) + } + } + } + if len(cols) == 0 && len(rows) == 0 { + return nil + } + if len(rows) > maxGridRows { + rows = rows[:maxGridRows] + } + rc := len(rows) + switch v := m["rowCount"].(type) { + case int: + rc = v + case int64: + rc = int(v) + case float64: + rc = int(v) + } + sqlText, _ := args["sql"].(string) + if sqlText == "" { + sqlText, _ = args["command"].(string) + } + return &AgentQueryResult{ + Columns: cols, + Rows: rows, + RowCount: rc, + SQL: sqlText, + Tool: tool, + } +} + func summarizeToolResult(tool string, out any) string { m, ok := out.(map[string]any) if !ok { diff --git a/server/internal/ai/agent_stream.go b/server/internal/ai/agent_stream.go new file mode 100644 index 0000000..8c12ea7 --- /dev/null +++ b/server/internal/ai/agent_stream.go @@ -0,0 +1,79 @@ +package ai + +import ( + "encoding/json" + "fmt" + "net/http" +) + +// AgentEventType SSE / 回调事件类型。 +const ( + AgentEventStepStart = "step_start" + AgentEventStep = "step" + AgentEventDelta = "delta" + AgentEventDone = "done" + AgentEventError = "error" +) + +// AgentQueryResult 可供「在结果区打开」的行列结果。 +type AgentQueryResult struct { + Columns []string `json:"columns"` + Rows []map[string]interface{} `json:"rows"` + RowCount int `json:"rowCount"` + SQL string `json:"sql,omitempty"` + Tool string `json:"tool,omitempty"` +} + +// AgentEvent 流式推送的一项。 +type AgentEvent struct { + Type string `json:"type"` + Step *AgentStep `json:"step,omitempty"` + Delta string `json:"delta,omitempty"` + Content string `json:"content,omitempty"` + SQL string `json:"sql,omitempty"` + Steps []AgentStep `json:"steps,omitempty"` + Rounds int `json:"rounds,omitempty"` + Error string `json:"error,omitempty"` + Result *AgentQueryResult `json:"result,omitempty"` +} + +// AgentEmitter 可选；nil 时仅返回最终 AgentRes（非流式 fallback）。 +type AgentEmitter func(AgentEvent) + +func emitEvent(emit AgentEmitter, ev AgentEvent) { + if emit != nil { + emit(ev) + } +} + +func chunkDelta(s string, size int) []string { + if s == "" { + return nil + } + if size <= 0 { + size = 32 + } + r := []rune(s) + var out []string + for i := 0; i < len(r); i += size { + j := i + size + if j > len(r) { + j = len(r) + } + out = append(out, string(r[i:j])) + } + return out +} + +// WriteAgentSSE 写一条 SSE 事件并 Flush。 +func WriteAgentSSE(w http.ResponseWriter, flusher http.Flusher, ev AgentEvent) error { + b, err := json.Marshal(ev) + if err != nil { + return err + } + if _, err := fmt.Fprintf(w, "event: %s\ndata: %s\n\n", ev.Type, b); err != nil { + return err + } + flusher.Flush() + return nil +} diff --git a/server/internal/ai/agent_test.go b/server/internal/ai/agent_test.go index 6bb7290..699b3df 100644 --- a/server/internal/ai/agent_test.go +++ b/server/internal/ai/agent_test.go @@ -202,7 +202,7 @@ func TestParseToolsTurn(t *testing.T) { }, }}, }) - turn, err := parseToolsTurn(raw, 200) + turn, err := parseToolsTurnChat(raw, 200) if err != nil || len(turn.ToolCalls) != 1 || turn.ToolCalls[0].Function.Name != "opsxdb_list_tables" { t.Fatalf("parse: %+v %v", turn, err) } diff --git a/server/internal/ai/ai_test.go b/server/internal/ai/ai_test.go new file mode 100644 index 0000000..605dd27 --- /dev/null +++ b/server/internal/ai/ai_test.go @@ -0,0 +1,100 @@ +package ai + +import ( + "encoding/json" + "net/http" + "strings" + "testing" +) + +func TestValidateEndpoint(t *testing.T) { + if err := validateEndpoint("https://api.deepseek.com/v1", false); err != nil { + t.Fatalf("https 应直接放行: %v", err) + } + if err := validateEndpoint("http://127.0.0.1:11434/v1", false); err == nil { + t.Fatal("未开启 allowInsecure 时 http 应拒绝") + } + if err := validateEndpoint("http://127.0.0.1:11434/v1", true); err != nil { + t.Fatalf("allowInsecure 下本地 http 应放行: %v", err) + } + if err := validateEndpoint("ftp://example.com", true); err == nil { + t.Fatal("非 HTTP/HTTPS scheme 应拒绝") + } + if err := validateEndpoint("not a url", false); err == nil { + t.Fatal("非法 URL 应拒绝") + } +} + +func TestEndpointJoin(t *testing.T) { + if got := endpoint("https://api.example.com/v1/", "chat/completions"); got != "https://api.example.com/v1/chat/completions" { + t.Fatalf("拼接错误: %s", got) + } + if got := endpoint("https://api.example.com/v1/responses", "responses"); got != "https://api.example.com/v1/responses" { + t.Fatalf("已含 suffix 不应重复: %s", got) + } + if got := anthropicEndpoint("https://api.anthropic.com"); got != "https://api.anthropic.com/v1/messages" { + t.Fatalf("anthropic 无 /v1 应补: %s", got) + } + if got := anthropicEndpoint("https://api.anthropic.com/v1"); got != "https://api.anthropic.com/v1/messages" { + t.Fatalf("anthropic 已含 /v1 不应重复: %s", got) + } +} + +func TestNewRemoteProviderInfersProtocolAndAuth(t *testing.T) { + legacy := NewRemoteProvider(Config{Provider: "openai", BaseURL: "https://api.openai.com/v1"}) + if legacy.protocol != ProtoChat || legacy.authStyle != AuthBearer { + t.Fatalf("旧配置应推断 openai-chat/bearer, got %s/%s", legacy.protocol, legacy.authStyle) + } + anthropic := NewRemoteProvider(Config{Provider: "anthropic", BaseURL: "https://api.anthropic.com"}) + if anthropic.protocol != ProtoAnthropic || anthropic.authStyle != AuthXAPIKey { + t.Fatalf("anthropic 预设应推断原生协议/x-api-key, got %s/%s", anthropic.protocol, anthropic.authStyle) + } + local := NewRemoteProvider(Config{Provider: "ollama", Protocol: ProtoChat, AuthStyle: AuthNone, BaseURL: "http://127.0.0.1:11434/v1", AllowInsecureHTTP: true}) + if local.authStyle != AuthNone { + t.Fatalf("本地免鉴权应保留 none, got %s", local.authStyle) + } +} + +func TestParseContentByProtocol(t *testing.T) { + chatRaw, _ := json.Marshal(map[string]any{ + "choices": []map[string]any{{"message": map[string]string{"role": "assistant", "content": " ```sql\nSELECT 1;\n``` "}}}, + }) + got, err := parseContent(ProtoChat, chatRaw, http.StatusOK) + if err != nil || strings.TrimSpace(got) != "```sql\nSELECT 1;\n```" { + t.Fatalf("chat 解析失败: %q, %v", got, err) + } + if s := extractSQL(got); s != "SELECT 1;" { + t.Fatalf("sql 代码块提取错误: %q", s) + } + + anthropicRaw, _ := json.Marshal(map[string]any{ + "content": []map[string]any{{"type": "text", "text": "你好"}, {"type": "tool_use", "text": "忽略"}}, + }) + got, err = parseContent(ProtoAnthropic, anthropicRaw, http.StatusOK) + if err != nil || got != "你好" { + t.Fatalf("anthropic 解析失败: %q, %v", got, err) + } + + responsesRaw, _ := json.Marshal(map[string]any{ + "output": []map[string]any{{"content": []map[string]string{{"text": "SELECT "}, {"text": "2;"}}}}, + }) + got, err = parseContent(ProtoResponses, responsesRaw, http.StatusOK) + if err != nil || got != "SELECT 2;" { + t.Fatalf("responses 解析失败: %q, %v", got, err) + } + + errRaw, _ := json.Marshal(map[string]any{"error": map[string]string{"message": "invalid api key"}}) + if _, err := parseContent(ProtoChat, errRaw, http.StatusUnauthorized); err == nil || !strings.Contains(err.Error(), "invalid api key") { + t.Fatalf("非 2xx 应提取服务端错误信息: %v", err) + } +} + +func TestBuildMessagesCoversActions(t *testing.T) { + msgs := BuildMessages(ChatReq{Action: "optimize", SQL: "SELECT * FROM t", Driver: "postgresql", Database: "db"}) + if len(msgs) != 2 || msgs[0].Role != "system" || msgs[1].Role != "user" { + t.Fatalf("应返回 [system, user], got %+v", msgs) + } + if !strings.Contains(msgs[1].Content, "SELECT * FROM t") { + t.Fatal("user 消息应包含 SQL") + } +} diff --git a/server/internal/ai/config.go b/server/internal/ai/config.go index dab3d60..7656d1d 100644 --- a/server/internal/ai/config.go +++ b/server/internal/ai/config.go @@ -7,12 +7,16 @@ import ( "strings" ) -// Config AI 配置;provider: stub|openai(OpenAI 兼容协议)。 +// Config AI 配置;provider 为预设 id(stub|openai|deepseek|anthropic|ollama|custom…), +// protocol/authStyle/allowInsecureHTTP 决定远端调用行为(见 remote.go)。 type Config struct { - Provider string `json:"provider"` - BaseURL string `json:"baseURL"` - APIKey string `json:"apiKey,omitempty"` - Model string `json:"model"` + Provider string `json:"provider"` + BaseURL string `json:"baseURL"` + APIKey string `json:"apiKey,omitempty"` + Model string `json:"model"` + Protocol string `json:"protocol,omitempty"` + AuthStyle string `json:"authStyle,omitempty"` + AllowInsecureHTTP bool `json:"allowInsecureHttp,omitempty"` } func configPath() (string, error) { @@ -47,6 +51,13 @@ func LoadConfig() Config { if f.Model != "" { cfg.Model = f.Model } + if f.Protocol != "" { + cfg.Protocol = f.Protocol + } + if f.AuthStyle != "" { + cfg.AuthStyle = f.AuthStyle + } + cfg.AllowInsecureHTTP = f.AllowInsecureHTTP } } } diff --git a/server/internal/ai/openai.go b/server/internal/ai/openai.go deleted file mode 100644 index 7e1766a..0000000 --- a/server/internal/ai/openai.go +++ /dev/null @@ -1,117 +0,0 @@ -package ai - -import ( - "bytes" - "context" - "encoding/json" - "fmt" - "io" - "net/http" - "regexp" - "strings" - "time" -) - -const ( - defaultOpenAIBaseURL = "https://api.openai.com/v1" - defaultOpenAIModel = "gpt-4o-mini" -) - -// OpenAIProvider OpenAI 兼容协议实现(DeepSeek/通义/本地 vLLM 同协议可用)。 -type OpenAIProvider struct { - baseURL string - apiKey string - model string - hc *http.Client -} - -func NewOpenAIProvider(cfg Config) *OpenAIProvider { - base := strings.TrimRight(strings.TrimSpace(cfg.BaseURL), "/") - if base == "" { - base = defaultOpenAIBaseURL - } - model := strings.TrimSpace(cfg.Model) - if model == "" { - model = defaultOpenAIModel - } - return &OpenAIProvider{ - baseURL: base, - apiKey: cfg.APIKey, - model: model, - hc: &http.Client{Timeout: 60 * time.Second}, - } -} - -func (p *OpenAIProvider) Name() string { return "openai" } - -type chatCompletionReq struct { - Model string `json:"model"` - Messages []Message `json:"messages"` -} - -type chatCompletionRes struct { - Choices []struct { - Message Message `json:"message"` - } `json:"choices"` - Error *struct { - Message string `json:"message"` - } `json:"error"` -} - -// Complete 单次非流式调用;ctx 全透传,HTTP 层 60s 超时兜底。 -func (p *OpenAIProvider) Complete(ctx context.Context, req ChatReq) (ChatRes, error) { - model := p.model - if strings.TrimSpace(req.Model) != "" { - model = strings.TrimSpace(req.Model) - } - payload := chatCompletionReq{ - Model: model, - Messages: BuildMessages(req), - } - body, err := json.Marshal(payload) - if err != nil { - return ChatRes{}, err - } - httpReq, err := http.NewRequestWithContext(ctx, http.MethodPost, p.baseURL+"/chat/completions", bytes.NewReader(body)) - if err != nil { - return ChatRes{}, err - } - httpReq.Header.Set("Content-Type", "application/json") - httpReq.Header.Set("Authorization", "Bearer "+p.apiKey) - - resp, err := p.hc.Do(httpReq) - if err != nil { - return ChatRes{}, fmt.Errorf("AI 请求失败: %w", err) - } - defer resp.Body.Close() - raw, err := io.ReadAll(io.LimitReader(resp.Body, 4<<20)) - if err != nil { - return ChatRes{}, fmt.Errorf("读取 AI 响应失败: %w", err) - } - var out chatCompletionRes - if err := json.Unmarshal(raw, &out); err != nil { - return ChatRes{}, fmt.Errorf("解析 AI 响应失败(HTTP %d)", resp.StatusCode) - } - if out.Error != nil { - return ChatRes{}, fmt.Errorf("AI 服务返回错误: %s", out.Error.Message) - } - if resp.StatusCode != http.StatusOK { - return ChatRes{}, fmt.Errorf("AI 服务 HTTP %d", resp.StatusCode) - } - if len(out.Choices) == 0 { - return ChatRes{}, fmt.Errorf("AI 服务未返回候选结果") - } - content := strings.TrimSpace(out.Choices[0].Message.Content) - return ChatRes{Content: content, SQL: extractSQL(content)}, nil -} - -var sqlFence = regexp.MustCompile("(?s)```sql\\s*\n(.+?)```") - -// extractSQL 从模型回复中提取 ```sql 代码块内容。 -func extractSQL(content string) string { - m := sqlFence.FindStringSubmatch(content) - if len(m) < 2 { - return "" - } - return strings.TrimSpace(m[1]) -} diff --git a/server/internal/ai/remote.go b/server/internal/ai/remote.go new file mode 100644 index 0000000..8c0423f --- /dev/null +++ b/server/internal/ai/remote.go @@ -0,0 +1,292 @@ +package ai + +import ( + "bytes" + "context" + "encoding/json" + "fmt" + "io" + "net/http" + "net/url" + "regexp" + "strings" + "time" +) + +const ( + defaultOpenAIBaseURL = "https://api.openai.com/v1" + defaultOpenAIModel = "gpt-4o-mini" + remoteMaxTokens = 2048 +) + +// 协议:openai-chat 覆盖绝大多数兼容服务;openai-responses 为 OpenAI 官方;anthropic 为 Claude。 +const ( + ProtoChat = "openai-chat" + ProtoResponses = "openai-responses" + ProtoAnthropic = "anthropic" +) + +// 鉴权风格:bearer = Authorization: Bearer;x-api-key = Anthropic 风格头;none = 本地服务免鉴权。 +const ( + AuthBearer = "bearer" + AuthXAPIKey = "x-api-key" + AuthNone = "none" +) + +// RemoteProvider 远端 HTTP 模型实现,按协议分派 chat/responses/anthropic(对齐 opsxshell AiService)。 +type RemoteProvider struct { + protocol string + authStyle string + baseURL string + apiKey string + model string + allowHTTP bool + hc *http.Client +} + +// NewRemoteProvider 按配置构造;protocol 缺省按 provider 推断(anthropic → 原生协议)。 +func NewRemoteProvider(cfg Config) *RemoteProvider { + protocol := strings.TrimSpace(cfg.Protocol) + if protocol == "" { + protocol = ProtoChat + if strings.EqualFold(cfg.Provider, ProtoAnthropic) { + protocol = ProtoAnthropic + } + } + authStyle := strings.TrimSpace(cfg.AuthStyle) + if authStyle == "" { + authStyle = AuthBearer + if protocol == ProtoAnthropic { + authStyle = AuthXAPIKey + } + } + base := strings.TrimRight(strings.TrimSpace(cfg.BaseURL), "/") + if base == "" { + base = defaultOpenAIBaseURL + } + model := strings.TrimSpace(cfg.Model) + if model == "" { + model = defaultOpenAIModel + } + return &RemoteProvider{ + protocol: protocol, + authStyle: authStyle, + baseURL: base, + apiKey: cfg.APIKey, + model: model, + allowHTTP: cfg.AllowInsecureHTTP, + hc: &http.Client{Timeout: 120 * time.Second}, + } +} + +func (p *RemoteProvider) Name() string { return p.protocol } + +// validateEndpoint 默认只放行 https;本地/内网服务(Ollama、vLLM)需显式 AllowInsecureHTTP。 +func validateEndpoint(baseURL string, allowInsecure bool) error { + u, err := url.Parse(strings.TrimSpace(baseURL)) + if err != nil || u.Host == "" { + return fmt.Errorf("Base URL 格式无效") + } + switch u.Scheme { + case "https": + return nil + case "http": + if allowInsecure { + return nil + } + return fmt.Errorf("非 HTTPS Base URL 已拒绝;本地/内网服务请在设置中开启「允许 HTTP」") + default: + return fmt.Errorf("Base URL 仅支持 HTTP/HTTPS") + } +} + +// endpoint 拼接 basePath + suffix,已含 suffix 时不重复。 +func endpoint(base, suffix string) string { + b := strings.TrimRight(strings.TrimSpace(base), "/") + s := strings.TrimLeft(suffix, "/") + if strings.HasSuffix(b, "/"+s) { + return b + } + return b + "/" + s +} + +// anthropicEndpoint base 未含 /v1 时补 /v1/messages,否则直接 /messages。 +func anthropicEndpoint(base string) string { + b := strings.TrimRight(strings.TrimSpace(base), "/") + if strings.HasSuffix(b, "/v1") { + return endpoint(b, "messages") + } + return endpoint(b, "v1/messages") +} + +// Complete 单次非流式调用;req.Model 可临时覆盖配置模型。 +func (p *RemoteProvider) Complete(ctx context.Context, req ChatReq) (ChatRes, error) { + if err := validateEndpoint(p.baseURL, p.allowHTTP); err != nil { + return ChatRes{}, err + } + model := p.model + if strings.TrimSpace(req.Model) != "" { + model = strings.TrimSpace(req.Model) + } + msgs := BuildMessages(req) + var ( + urlStr string + payload any + ) + switch p.protocol { + case ProtoAnthropic: + urlStr = anthropicEndpoint(p.baseURL) + payload = map[string]any{ + "model": model, + "system": msgs[0].Content, + "messages": []Message{{Role: "user", Content: msgs[1].Content}}, + "max_tokens": remoteMaxTokens, + "temperature": 0.2, + } + case ProtoResponses: + urlStr = endpoint(p.baseURL, "responses") + payload = map[string]any{ + "model": model, + "instructions": msgs[0].Content, + "input": msgs[1].Content, + "max_output_tokens": remoteMaxTokens, + } + default: // openai-chat 及未知值按兼容协议处理 + urlStr = endpoint(p.baseURL, "chat/completions") + payload = map[string]any{ + "model": model, + "messages": msgs, + "max_tokens": remoteMaxTokens, + "temperature": 0.2, + } + } + body, err := json.Marshal(payload) + if err != nil { + return ChatRes{}, err + } + httpReq, err := http.NewRequestWithContext(ctx, http.MethodPost, urlStr, bytes.NewReader(body)) + if err != nil { + return ChatRes{}, err + } + httpReq.Header.Set("Content-Type", "application/json") + switch p.authStyle { + case AuthXAPIKey: + httpReq.Header.Set("x-api-key", p.apiKey) + httpReq.Header.Set("anthropic-version", "2023-06-01") + case AuthNone: // 本地免鉴权服务 + default: + httpReq.Header.Set("Authorization", "Bearer "+p.apiKey) + } + + resp, err := p.hc.Do(httpReq) + if err != nil { + return ChatRes{}, fmt.Errorf("AI 请求失败: %w", err) + } + defer resp.Body.Close() + raw, err := io.ReadAll(io.LimitReader(resp.Body, 4<<20)) + if err != nil { + return ChatRes{}, fmt.Errorf("读取 AI 响应失败: %w", err) + } + content, err := parseContent(p.protocol, raw, resp.StatusCode) + if err != nil { + return ChatRes{}, err + } + return ChatRes{Content: content, SQL: extractSQL(content)}, nil +} + +// parseContent 按协议解析响应体;非 2xx 时尽量提取服务端错误信息。 +func parseContent(protocol string, raw []byte, status int) (string, error) { + var errBody struct { + Error *struct { + Message string `json:"message"` + } `json:"error"` + Message string `json:"message"` + } + _ = json.Unmarshal(raw, &errBody) + if status < 200 || status >= 300 { + msg := "未知错误" + if errBody.Error != nil && errBody.Error.Message != "" { + msg = errBody.Error.Message + } else if errBody.Message != "" { + msg = errBody.Message + } + return "", fmt.Errorf("AI 服务 HTTP %d: %s", status, msg) + } + switch protocol { + case ProtoAnthropic: + var out struct { + Content []struct { + Type string `json:"type"` + Text string `json:"text"` + } `json:"content"` + } + if err := json.Unmarshal(raw, &out); err != nil { + return "", fmt.Errorf("解析 AI 响应失败(HTTP %d)", status) + } + var sb strings.Builder + for _, c := range out.Content { + if c.Type == "text" { + sb.WriteString(c.Text) + } + } + if sb.Len() == 0 { + return "", fmt.Errorf("AI 服务未返回文本内容") + } + return sb.String(), nil + case ProtoResponses: + var out struct { + OutputText string `json:"output_text"` + Output []struct { + Content []struct { + Text string `json:"text"` + } `json:"content"` + } `json:"output"` + } + if err := json.Unmarshal(raw, &out); err != nil { + return "", fmt.Errorf("解析 AI 响应失败(HTTP %d)", status) + } + if out.OutputText != "" { + return out.OutputText, nil + } + var sb strings.Builder + for _, o := range out.Output { + for _, c := range o.Content { + sb.WriteString(c.Text) + } + } + if sb.Len() == 0 { + return "", fmt.Errorf("AI 服务未返回文本内容") + } + return sb.String(), nil + default: + var out struct { + Choices []struct { + Message Message `json:"message"` + } `json:"choices"` + Error *struct { + Message string `json:"message"` + } `json:"error"` + } + if err := json.Unmarshal(raw, &out); err != nil { + return "", fmt.Errorf("解析 AI 响应失败(HTTP %d)", status) + } + if out.Error != nil { + return "", fmt.Errorf("AI 服务返回错误: %s", out.Error.Message) + } + if len(out.Choices) == 0 { + return "", fmt.Errorf("AI 服务未返回候选结果") + } + return strings.TrimSpace(out.Choices[0].Message.Content), nil + } +} + +var sqlFence = regexp.MustCompile("(?s)```sql\\s*\n(.+?)```") + +// extractSQL 从模型回复中提取 ```sql 代码块内容。 +func extractSQL(content string) string { + m := sqlFence.FindStringSubmatch(content) + if len(m) < 2 { + return "" + } + return strings.TrimSpace(m[1]) +} diff --git a/server/internal/ai/remote_tools.go b/server/internal/ai/remote_tools.go index aa6d170..07a3bca 100644 --- a/server/internal/ai/remote_tools.go +++ b/server/internal/ai/remote_tools.go @@ -24,35 +24,33 @@ type ToolChat interface { ChatWithTools(ctx context.Context, model string, messages []Message, tools []aitools.OpenAITool) (ToolsTurn, error) } -// ChatWithTools openai-chat 兼容协议的 tools 调用（DeepSeek / OpenAI / 多数兼容网关）。 -// anthropic / responses 本轮不支持 Agent tool calling，返回明确错误。 +// ChatWithTools 按协议分派：openai-chat / anthropic messages / openai-responses。 func (p *RemoteProvider) ChatWithTools(ctx context.Context, model string, messages []Message, tools []aitools.OpenAITool) (ToolsTurn, error) { if err := validateEndpoint(p.baseURL, p.allowHTTP); err != nil { return ToolsTurn{}, err } - if p.protocol != "" && p.protocol != ProtoChat { - return ToolsTurn{}, fmt.Errorf("Agent 模式目前仅支持 openai-chat 协议（当前: %s）；请改用 DeepSeek / OpenAI 兼容端点", p.protocol) - } m := strings.TrimSpace(model) if m == "" { m = p.model } - urlStr := endpoint(p.baseURL, "chat/completions") - payload := map[string]any{ - "model": m, - "messages": messages, - "tools": tools, - "tool_choice": "auto", - "max_tokens": remoteMaxTokens, - "temperature": 0.2, + switch p.protocol { + case ProtoAnthropic: + return p.chatToolsAnthropic(ctx, m, messages, tools) + case ProtoResponses: + return p.chatToolsResponses(ctx, m, messages, tools) + default: + return p.chatToolsOpenAI(ctx, m, messages, tools) } +} + +func (p *RemoteProvider) doJSON(ctx context.Context, urlStr string, payload any) ([]byte, int, error) { body, err := json.Marshal(payload) if err != nil { - return ToolsTurn{}, err + return nil, 0, err } httpReq, err := http.NewRequestWithContext(ctx, http.MethodPost, urlStr, bytes.NewReader(body)) if err != nil { - return ToolsTurn{}, err + return nil, 0, err } httpReq.Header.Set("Content-Type", "application/json") switch p.authStyle { @@ -65,17 +63,174 @@ func (p *RemoteProvider) ChatWithTools(ctx context.Context, model string, messag } resp, err := p.hc.Do(httpReq) if err != nil { - return ToolsTurn{}, fmt.Errorf("AI 请求失败: %w", err) + return nil, 0, fmt.Errorf("AI 请求失败: %w", err) } defer resp.Body.Close() raw, err := io.ReadAll(io.LimitReader(resp.Body, 4<<20)) if err != nil { - return ToolsTurn{}, fmt.Errorf("读取 AI 响应失败: %w", err) + return nil, resp.StatusCode, fmt.Errorf("读取 AI 响应失败: %w", err) } - return parseToolsTurn(raw, resp.StatusCode) + return raw, resp.StatusCode, nil } -func parseToolsTurn(raw []byte, status int) (ToolsTurn, error) { +func (p *RemoteProvider) chatToolsOpenAI(ctx context.Context, model string, messages []Message, tools []aitools.OpenAITool) (ToolsTurn, error) { + urlStr := endpoint(p.baseURL, "chat/completions") + payload := map[string]any{ + "model": model, + "messages": messages, + "tools": tools, + "tool_choice": "auto", + "max_tokens": remoteMaxTokens, + "temperature": 0.2, + } + raw, status, err := p.doJSON(ctx, urlStr, payload) + if err != nil { + return ToolsTurn{}, err + } + return parseToolsTurnChat(raw, status) +} + +func (p *RemoteProvider) chatToolsAnthropic(ctx context.Context, model string, messages []Message, tools []aitools.OpenAITool) (ToolsTurn, error) { + urlStr := anthropicEndpoint(p.baseURL) + system, anthMsgs := splitAnthropicMessages(messages) + payload := map[string]any{ + "model": model, + "max_tokens": remoteMaxTokens, + "temperature": 0.2, + "tools": openaiToolsToAnthropic(tools), + "messages": anthMsgs, + } + if system != "" { + payload["system"] = system + } + raw, status, err := p.doJSON(ctx, urlStr, payload) + if err != nil { + return ToolsTurn{}, err + } + return parseToolsTurnAnthropic(raw, status) +} + +func (p *RemoteProvider) chatToolsResponses(ctx context.Context, model string, messages []Message, tools []aitools.OpenAITool) (ToolsTurn, error) { + urlStr := endpoint(p.baseURL, "responses") + instructions, input := splitResponsesInput(messages) + payload := map[string]any{ + "model": model, + "instructions": instructions, + "input": input, + "tools": openaiToolsToResponses(tools), + "max_output_tokens": remoteMaxTokens, + } + raw, status, err := p.doJSON(ctx, urlStr, payload) + if err != nil { + return ToolsTurn{}, err + } + return parseToolsTurnResponses(raw, status) +} + +func splitAnthropicMessages(messages []Message) (system string, out []map[string]any) { + for _, m := range messages { + switch m.Role { + case "system": + if system == "" { + system = m.Content + } else { + system += "\n" + m.Content + } + case "assistant": + content := []map[string]any{} + if strings.TrimSpace(m.Content) != "" { + content = append(content, map[string]any{"type": "text", "text": m.Content}) + } + for _, tc := range m.ToolCalls { + var input any = map[string]any{} + _ = json.Unmarshal([]byte(tc.Function.Arguments), &input) + content = append(content, map[string]any{ + "type": "tool_use", + "id": tc.ID, + "name": tc.Function.Name, + "input": input, + }) + } + out = append(out, map[string]any{"role": "assistant", "content": content}) + case "tool": + out = append(out, map[string]any{ + "role": "user", + "content": []map[string]any{{ + "type": "tool_result", + "tool_use_id": m.ToolCallID, + "content": m.Content, + }}, + }) + default: // user + out = append(out, map[string]any{"role": "user", "content": m.Content}) + } + } + return system, out +} + +func splitResponsesInput(messages []Message) (instructions string, input []map[string]any) { + for _, m := range messages { + switch m.Role { + case "system": + if instructions == "" { + instructions = m.Content + } else { + instructions += "\n" + m.Content + } + case "assistant": + item := map[string]any{"role": "assistant", "content": m.Content} + if len(m.ToolCalls) > 0 { + calls := make([]map[string]any, 0, len(m.ToolCalls)) + for _, tc := range m.ToolCalls { + calls = append(calls, map[string]any{ + "id": tc.ID, + "type": "function_call", + "name": tc.Function.Name, + "arguments": tc.Function.Arguments, + }) + } + item["tool_calls"] = calls + } + input = append(input, item) + case "tool": + input = append(input, map[string]any{ + "type": "function_call_output", + "call_id": m.ToolCallID, + "output": m.Content, + }) + default: + input = append(input, map[string]any{"role": "user", "content": m.Content}) + } + } + return instructions, input +} + +func openaiToolsToAnthropic(tools []aitools.OpenAITool) []map[string]any { + out := make([]map[string]any, 0, len(tools)) + for _, t := range tools { + out = append(out, map[string]any{ + "name": t.Function.Name, + "description": t.Function.Description, + "input_schema": t.Function.Parameters, + }) + } + return out +} + +func openaiToolsToResponses(tools []aitools.OpenAITool) []map[string]any { + out := make([]map[string]any, 0, len(tools)) + for _, t := range tools { + out = append(out, map[string]any{ + "type": "function", + "name": t.Function.Name, + "description": t.Function.Description, + "parameters": t.Function.Parameters, + }) + } + return out +} + +func httpErrMessage(raw []byte, status int) error { var errBody struct { Error *struct { Message string `json:"message"` @@ -90,7 +245,14 @@ func parseToolsTurn(raw []byte, status int) (ToolsTurn, error) { } else if errBody.Message != "" { msg = errBody.Message } - return ToolsTurn{}, fmt.Errorf("AI 服务 HTTP %d: %s", status, msg) + return fmt.Errorf("AI 服务 HTTP %d: %s", status, msg) + } + return nil +} + +func parseToolsTurnChat(raw []byte, status int) (ToolsTurn, error) { + if err := httpErrMessage(raw, status); err != nil { + return ToolsTurn{}, err } var out struct { Choices []struct { @@ -119,25 +281,108 @@ func parseToolsTurn(raw []byte, status int) (ToolsTurn, error) { if ch.Message.Content != nil { content = strings.TrimSpace(*ch.Message.Content) } - return ToolsTurn{ - Content: content, - ToolCalls: ch.Message.ToolCalls, - RawFinish: ch.FinishReason, - }, nil + return ToolsTurn{Content: content, ToolCalls: ch.Message.ToolCalls, RawFinish: ch.FinishReason}, nil +} + +func parseToolsTurnAnthropic(raw []byte, status int) (ToolsTurn, error) { + if err := httpErrMessage(raw, status); err != nil { + return ToolsTurn{}, err + } + var out struct { + Content []struct { + Type string `json:"type"` + Text string `json:"text"` + ID string `json:"id"` + Name string `json:"name"` + Input json.RawMessage `json:"input"` + } `json:"content"` + StopReason string `json:"stop_reason"` + } + if err := json.Unmarshal(raw, &out); err != nil { + return ToolsTurn{}, fmt.Errorf("解析 Anthropic 响应失败(HTTP %d)", status) + } + var content strings.Builder + var calls []ToolCall + for _, c := range out.Content { + switch c.Type { + case "text": + content.WriteString(c.Text) + case "tool_use": + args := "{}" + if len(c.Input) > 0 { + args = string(c.Input) + } + calls = append(calls, ToolCall{ + ID: c.ID, + Type: "function", + Function: ToolCallFunc{Name: c.Name, Arguments: args}, + }) + } + } + return ToolsTurn{Content: strings.TrimSpace(content.String()), ToolCalls: calls, RawFinish: out.StopReason}, nil +} + +func parseToolsTurnResponses(raw []byte, status int) (ToolsTurn, error) { + if err := httpErrMessage(raw, status); err != nil { + return ToolsTurn{}, err + } + var out struct { + OutputText string `json:"output_text"` + Output []struct { + Type string `json:"type"` + ID string `json:"id"` + CallID string `json:"call_id"` + Name string `json:"name"` + Arguments string `json:"arguments"` + Content []struct { + Type string `json:"type"` + Text string `json:"text"` + } `json:"content"` + } `json:"output"` + } + if err := json.Unmarshal(raw, &out); err != nil { + return ToolsTurn{}, fmt.Errorf("解析 responses 响应失败(HTTP %d)", status) + } + var content strings.Builder + if out.OutputText != "" { + content.WriteString(out.OutputText) + } + var calls []ToolCall + for _, o := range out.Output { + switch o.Type { + case "message": + for _, c := range o.Content { + if c.Type == "output_text" || c.Type == "text" || c.Text != "" { + content.WriteString(c.Text) + } + } + case "function_call": + id := o.CallID + if id == "" { + id = o.ID + } + args := o.Arguments + if args == "" { + args = "{}" + } + calls = append(calls, ToolCall{ + ID: id, + Type: "function", + Function: ToolCallFunc{Name: o.Name, Arguments: args}, + }) + } + } + return ToolsTurn{Content: strings.TrimSpace(content.String()), ToolCalls: calls}, nil } // DefaultToolChat 返回可跑 Agent 的远端 Provider；stub / 缺 key 时返回错误说明。 func DefaultToolChat() (ToolChat, error) { cfg := LoadConfig() if cfg.Provider == "" || cfg.Provider == "stub" { - return nil, fmt.Errorf("Agent 需要真实 LLM：请在设置中配置 DeepSeek / OpenAI 等（stub 仅支持 Ask 生成）") + return nil, fmt.Errorf("Agent 需要真实 LLM：请在设置中配置 DeepSeek / OpenAI / Anthropic 等（stub 仅支持 Ask 生成）") } if cfg.APIKey == "" && cfg.AuthStyle != AuthNone { return nil, fmt.Errorf("Agent 需要 API Key：请在设置中填写后重试") } - p := NewRemoteProvider(cfg) - if p.protocol != "" && p.protocol != ProtoChat { - return nil, fmt.Errorf("Agent 模式目前仅支持 openai-chat 协议（当前: %s）", p.protocol) - } - return p, nil + return NewRemoteProvider(cfg), nil } diff --git a/server/internal/ai/remote_tools_test.go b/server/internal/ai/remote_tools_test.go new file mode 100644 index 0000000..abde9a0 --- /dev/null +++ b/server/internal/ai/remote_tools_test.go @@ -0,0 +1,72 @@ +package ai + +import ( + "encoding/json" + "strings" + "testing" + + "opsxdb/server/internal/aitools" +) + +func TestParseToolsTurnAnthropic(t *testing.T) { + raw, _ := json.Marshal(map[string]any{ + "stop_reason": "tool_use", + "content": []map[string]any{ + {"type": "text", "text": "先列表"}, + {"type": "tool_use", "id": "tu_1", "name": "opsxdb_list_tables", "input": map[string]any{"connection_id": "c1"}}, + }, + }) + turn, err := parseToolsTurnAnthropic(raw, 200) + if err != nil { + t.Fatal(err) + } + if turn.Content != "先列表" { + t.Fatalf("content=%q", turn.Content) + } + if len(turn.ToolCalls) != 1 || turn.ToolCalls[0].Function.Name != "opsxdb_list_tables" { + t.Fatalf("calls=%+v", turn.ToolCalls) + } + if !strings.Contains(turn.ToolCalls[0].Function.Arguments, "c1") { + t.Fatalf("args=%s", turn.ToolCalls[0].Function.Arguments) + } +} + +func TestParseToolsTurnResponses(t *testing.T) { + raw, _ := json.Marshal(map[string]any{ + "output": []map[string]any{ + {"type": "function_call", "call_id": "call_9", "name": "opsxdb_execute_query", "arguments": `{"connection_id":"c1","sql":"SELECT 1"}`}, + {"type": "message", "content": []map[string]any{{"type": "output_text", "text": "done"}}}, + }, + }) + turn, err := parseToolsTurnResponses(raw, 200) + if err != nil { + t.Fatal(err) + } + if len(turn.ToolCalls) != 1 || turn.ToolCalls[0].ID != "call_9" { + t.Fatalf("calls=%+v", turn.ToolCalls) + } + if turn.Content != "done" { + t.Fatalf("content=%q", turn.Content) + } +} + +func TestOpenAIToolsToAnthropic(t *testing.T) { + tools := aitools.OpenAITools() + anth := openaiToolsToAnthropic(tools) + if len(anth) != len(tools) { + t.Fatalf("%d vs %d", len(anth), len(tools)) + } + if anth[0]["name"] != tools[0].Function.Name { + t.Fatal(anth[0]["name"]) + } + if anth[0]["input_schema"] == nil { + t.Fatal("missing input_schema") + } +} + +func TestChunkDelta(t *testing.T) { + parts := chunkDelta("abcdefghij", 3) + if len(parts) != 4 { + t.Fatalf("%v", parts) + } +} diff --git a/server/internal/aitools/audit.go b/server/internal/aitools/audit.go index 47026de..1674599 100644 --- a/server/internal/aitools/audit.go +++ b/server/internal/aitools/audit.go @@ -14,6 +14,7 @@ type AuditEntry struct { Tool string `json:"tool"` ConnectionID string `json:"connection_id,omitempty"` Mode string `json:"mode"` + Source string `json:"source,omitempty"` // mcp | agent OK bool `json:"ok"` Summary string `json:"summary,omitempty"` Error string `json:"error,omitempty"` diff --git a/server/internal/aitools/policy.go b/server/internal/aitools/policy.go index 9e8225e..37095ef 100644 --- a/server/internal/aitools/policy.go +++ b/server/internal/aitools/policy.go @@ -81,10 +81,14 @@ func LoadPolicy() (Policy, error) { } func (p Policy) IsReadOnly() bool { - // P0：即使配置写了 read_write，执行层仍按只读拦截；此 helper 表示「策略意图」。 return p.Mode != ModeReadWrite } +// WritesEnabled P1：仅当 mcp.json mode=read_write 时允许写路径（仍须确认令牌）。 +func (p Policy) WritesEnabled() bool { + return strings.EqualFold(strings.TrimSpace(p.Mode), ModeReadWrite) +} + func (p Policy) AllowConnection(id string) error { id = strings.TrimSpace(id) if id == "" { diff --git a/server/internal/aitools/runtime.go b/server/internal/aitools/runtime.go index d6d0da3..1086fe6 100644 --- a/server/internal/aitools/runtime.go +++ b/server/internal/aitools/runtime.go @@ -46,6 +46,13 @@ type Runtime struct { ListTbl func(ctx context.Context, c db.ConnReq) ([]string, error) ListCol func(ctx context.Context, c db.ConnReq) ([]db.ColumnInfo, error) Audit func(AuditEntry) + // RejectCallerConfirm 为 true 时（Agent 路径），忽略调用方 confirm_danger， + // 写操作只返回 needConfirm，由用户在 UI 显式确认后再走 /api/ai/agent/confirm。 + RejectCallerConfirm bool + // AllowConfirmedWrite 为 true 时，已校验的确认令牌可放行写（仅 confirm 端点设置）。 + AllowConfirmedWrite bool + // Source 写入审计：mcp | agent。 + Source string } func NewRuntime(policy Policy) *Runtime { @@ -67,6 +74,9 @@ func (r *Runtime) audit(e AuditEntry) { e.Mode = ModeReadOnly } } + if e.Source == "" { + e.Source = r.Source + } if r.Audit != nil { r.Audit(e) return @@ -213,6 +223,29 @@ func argInt(args map[string]any, key string, def int) int { } } +func argBool(args map[string]any, key string) bool { + if args == nil { + return false + } + v, ok := args[key] + if !ok || v == nil { + return false + } + switch t := v.(type) { + case bool: + return t + case string: + return strings.EqualFold(t, "true") || t == "1" + case float64: + return t != 0 + case int: + return t != 0 + default: + return false + } +} + + func (r *Runtime) listDatabases(ctx context.Context, args map[string]any) (any, error) { id := argString(args, "connection_id") c, err := r.resolveAllowed(id) @@ -352,9 +385,27 @@ func (r *Runtime) executeQuery(ctx context.Context, args map[string]any) (any, e if sqlText == "" { return fail("sql required") } - // P0 纯只读：写/DDL 一律拒绝；不做 write mode / confirm token。 - if ok, reason := IsReadOnlySQL(sqlText); !ok { - return fail(reason) + callerConfirm := argBool(args, "confirm_danger") || argBool(args, "confirmDanger") + if r.RejectCallerConfirm { + callerConfirm = false + } + confirmed := r.AllowConfirmedWrite || callerConfirm + + readOK, readReason := IsReadOnlySQL(sqlText) + if !readOK { + if !r.Policy.WritesEnabled() { + return fail(readReason) + } + // P1：可写已开启，但无确认令牌则返回 needConfirm（永不静默写）。 + if !confirmed { + r.audit(AuditEntry{Tool: ToolExecuteQuery, ConnectionID: id, OK: false, Error: "needConfirm", Summary: truncate(sqlText, 80)}) + return map[string]any{ + "needConfirm": true, + "sql": sqlText, + "message": "write requires explicit user confirmation", + "reason": readReason, + }, nil + } } c, err := r.resolveAllowed(id) @@ -371,7 +422,7 @@ func (r *Runtime) executeQuery(ctx context.Context, args map[string]any) (any, e } req.SQL = sqlText req.Limit = limit - req.ConfirmDanger = false // 永不带确认放行 + req.ConfirmDanger = confirmed && !readOK res := r.execSQL(ctx, req) if res.NeedConfirm { @@ -407,11 +458,27 @@ func (r *Runtime) executeRedis(ctx context.Context, args map[string]any) (any, e if command == "" { return fail("command required") } - cmd, err := CheckRedisReadOnly(command) - if err != nil { - return fail(err.Error()) + callerConfirm := argBool(args, "confirm_danger") || argBool(args, "confirmDanger") + if r.RejectCallerConfirm { + callerConfirm = false + } + confirmed := r.AllowConfirmedWrite || callerConfirm + + cmd, rerr := CheckRedisReadOnly(command) + if rerr != nil { + if redisDeniedAlways[cmd] || !r.Policy.WritesEnabled() { + return fail(rerr.Error()) + } + if !confirmed { + r.audit(AuditEntry{Tool: ToolExecuteRedisCommand, ConnectionID: id, OK: false, Error: "needConfirm", Summary: truncate(command, 80)}) + return map[string]any{ + "needConfirm": true, + "command": command, + "message": "redis write requires explicit user confirmation", + "reason": rerr.Error(), + }, nil + } } - _ = cmd c, err := r.resolveAllowed(id) if err != nil { @@ -426,7 +493,7 @@ func (r *Runtime) executeRedis(ctx context.Context, args map[string]any) (any, e req.Database = database } req.SQL = command - req.ConfirmDanger = false + req.ConfirmDanger = confirmed && rerr != nil res := r.execSQL(ctx, req) if res.NeedConfirm { diff --git a/server/internal/aitools/runtime_test.go b/server/internal/aitools/runtime_test.go index 041e338..8b913c5 100644 --- a/server/internal/aitools/runtime_test.go +++ b/server/internal/aitools/runtime_test.go @@ -227,3 +227,75 @@ func TestToolNamesStable(t *testing.T) { t.Fatalf("%v", names) } } + +func TestExecuteQueryNeedConfirmInReadWrite(t *testing.T) { + called := false + rt := NewRuntime(Policy{ + Mode: ModeReadWrite, + AllowedConnectionIDs: []string{"c1"}, + DefaultLimit: 100, + MaxLimit: 500, + }) + rt.FindConn = func(id string) (*SavedConnection, error) { + return &SavedConnection{ID: id, Name: "d", Driver: "mysql", Host: "127.0.0.1", Port: 3306}, nil + } + rt.ExecSQL = func(_ context.Context, c db.ConnReq) db.ExecResult { + called = true + return db.ExecResult{RowCount: 1} + } + rt.Audit = func(AuditEntry) {} + + out, err := rt.Call(context.Background(), ToolExecuteQuery, map[string]any{ + "connection_id": "c1", + "sql": "DELETE FROM users WHERE id=1", + }) + if err != nil { + t.Fatal(err) + } + m := out.(map[string]any) + if m["needConfirm"] != true { + t.Fatalf("want needConfirm, got %#v", m) + } + if called { + t.Fatal("must not exec without confirm") + } + + out2, err := rt.Call(context.Background(), ToolExecuteQuery, map[string]any{ + "connection_id": "c1", + "sql": "DELETE FROM users WHERE id=1", + "confirm_danger": true, + }) + if err != nil { + t.Fatal(err) + } + if !called { + t.Fatal("expected exec after confirm") + } + _ = out2 +} + +func TestExecuteQueryRejectConfirmWhenAgentFlag(t *testing.T) { + called := false + rt := NewRuntime(Policy{Mode: ModeReadWrite, AllowedConnectionIDs: []string{"c1"}, DefaultLimit: 100, MaxLimit: 500}) + rt.RejectCallerConfirm = true + rt.FindConn = func(id string) (*SavedConnection, error) { + return &SavedConnection{ID: id, Driver: "mysql", Host: "h", Port: 1}, nil + } + rt.ExecSQL = func(_ context.Context, c db.ConnReq) db.ExecResult { + called = true + return db.ExecResult{} + } + rt.Audit = func(AuditEntry) {} + out, err := rt.Call(context.Background(), ToolExecuteQuery, map[string]any{ + "connection_id": "c1", "sql": "UPDATE t SET a=1", "confirm_danger": true, + }) + if err != nil { + t.Fatal(err) + } + if out.(map[string]any)["needConfirm"] != true { + t.Fatalf("%#v", out) + } + if called { + t.Fatal("agent must not auto-confirm") + } +} diff --git a/server/internal/aitools/schema.go b/server/internal/aitools/schema.go index b3ebcdb..5bf72df 100644 --- a/server/internal/aitools/schema.go +++ b/server/internal/aitools/schema.go @@ -52,22 +52,24 @@ func ToolSpecs() []ToolSpec { }, { Name: ToolExecuteQuery, - Description: "Execute a READ-ONLY SQL query (SELECT/SHOW/EXPLAIN/WITH…SELECT). Writes and DDL are rejected in P0. Default limit 100, max 500.", + Description: "Execute SQL. Default read-only (SELECT/SHOW/EXPLAIN/WITH…SELECT). Writes require mcp.json mode=read_write AND confirm_danger=true (never auto). Default limit 100, max 500.", Properties: map[string]any{ - "connection_id": connID, - "sql": map[string]any{"type": "string", "description": "Read-only SQL"}, - "database": database, - "limit": map[string]any{"type": "integer", "description": "Row limit (default 100, max 500)"}, + "connection_id": connID, + "sql": map[string]any{"type": "string", "description": "SQL statement"}, + "database": database, + "limit": map[string]any{"type": "integer", "description": "Row limit (default 100, max 500)"}, + "confirm_danger": map[string]any{"type": "boolean", "description": "Required for writes when mode=read_write; Agent UI sets this only after user confirms"}, }, Required: []string{"connection_id", "sql"}, }, { Name: ToolExecuteRedisCommand, - Description: "Execute a READ-ONLY Redis command (GET/SCAN/TYPE/TTL/…). FLUSH*/KEYS and writes are rejected in P0.", + Description: "Execute Redis command. Default read-only whitelist. Writes require mode=read_write AND confirm_danger=true. FLUSH*/KEYS/CONFIG always denied.", Properties: map[string]any{ - "connection_id": connID, - "command": map[string]any{"type": "string", "description": "Single Redis command line"}, - "database": database, + "connection_id": connID, + "command": map[string]any{"type": "string", "description": "Single Redis command line"}, + "database": database, + "confirm_danger": map[string]any{"type": "boolean", "description": "Required for writes when mode=read_write"}, }, Required: []string{"connection_id", "command"}, }, diff --git a/server/internal/aitools/status.go b/server/internal/aitools/status.go new file mode 100644 index 0000000..d63f23a --- /dev/null +++ b/server/internal/aitools/status.go @@ -0,0 +1,98 @@ +package aitools + +import ( + "bufio" + "encoding/json" + "os" + "time" +) + +const mcpActiveWindow = 5 * time.Minute + +// MCPStatus 供侧栏 MCP 状态点 / 设置页轻量指示。 +type MCPStatus struct { + Configured bool `json:"configured"` + Mode string `json:"mode"` + WhitelistCount int `json:"whitelistCount"` + WritesEnabled bool `json:"writesEnabled"` + Active bool `json:"active"` // 近 5 分钟有 source=mcp 审计 + Recent []AuditEntry `json:"recent,omitempty"` + PolicyPath string `json:"policyPath,omitempty"` + AuditPath string `json:"auditPath,omitempty"` +} + +// GetMCPStatus 读 mcp.json + 尾部审计；不探测外部 Cursor 进程。 +func GetMCPStatus() MCPStatus { + st := MCPStatus{Mode: ModeReadOnly} + if path, err := PolicyPath(); err == nil { + st.PolicyPath = path + if _, err := os.Stat(path); err == nil { + st.Configured = true + } + } + if p, err := LoadPolicy(); err == nil { + st.Mode = p.Mode + if st.Mode == "" { + st.Mode = ModeReadOnly + } + st.WhitelistCount = len(p.AllowedConnectionIDs) + st.WritesEnabled = p.WritesEnabled() + st.Configured = st.Configured || st.WhitelistCount > 0 + } + if ap, err := AuditPath(); err == nil { + st.AuditPath = ap + recent := tailAudit(ap, 20) + st.Recent = filterSource(recent, "mcp") + if len(st.Recent) == 0 { + // 兼容旧审计无 source：仍展示最近条目，但不标 active + st.Recent = recent + } else { + cutoff := time.Now().Add(-mcpActiveWindow) + for _, e := range st.Recent { + if t, err := time.Parse(time.RFC3339, e.Time); err == nil && t.After(cutoff) { + st.Active = true + break + } + } + } + } + return st +} + +func filterSource(in []AuditEntry, source string) []AuditEntry { + var out []AuditEntry + for _, e := range in { + if e.Source == source { + out = append(out, e) + } + } + return out +} + +func tailAudit(path string, n int) []AuditEntry { + f, err := os.Open(path) + if err != nil { + return nil + } + defer f.Close() + var lines []string + sc := bufio.NewScanner(f) + sc.Buffer(make([]byte, 0, 64*1024), 1024*1024) + for sc.Scan() { + lines = append(lines, sc.Text()) + if len(lines) > n*4 { // 粗裁 + lines = lines[len(lines)-n*2:] + } + } + if len(lines) > n { + lines = lines[len(lines)-n:] + } + out := make([]AuditEntry, 0, len(lines)) + for _, line := range lines { + var e AuditEntry + if json.Unmarshal([]byte(line), &e) == nil { + out = append(out, e) + } + } + return out +} diff --git a/server/internal/server/router.go b/server/internal/server/router.go index ba07e04..4524ef9 100644 --- a/server/internal/server/router.go +++ b/server/internal/server/router.go @@ -475,6 +475,7 @@ func NewRouter(token string) *gin.Engine { }) // 侧栏 Agent：多轮 tool call → 同一 aitools.Runtime（进程内，不绕 stdio）。Ask 仍走 /ai/chat。 + // stream=true 或 Accept: text/event-stream 时走 SSE（步骤边推 + 正文分块）；否则 JSON fallback。 api.POST("/ai/agent", func(c *gin.Context) { var req ai.AgentReq if err := c.ShouldBindJSON(&req); err != nil { @@ -491,6 +492,30 @@ func NewRouter(token string) *gin.Engine { c.JSON(http.StatusInternalServerError, gin.H{"error": err.Error()}) return } + rt.Source = "agent" + rt.RejectCallerConfirm = true + + wantStream := req.Stream || strings.Contains(c.GetHeader("Accept"), "text/event-stream") + if wantStream { + flusher, ok := c.Writer.(http.Flusher) + if !ok { + c.JSON(http.StatusInternalServerError, gin.H{"error": "streaming unsupported"}) + return + } + c.Header("Content-Type", "text/event-stream") + c.Header("Cache-Control", "no-cache") + c.Header("Connection", "keep-alive") + c.Status(http.StatusOK) + flusher.Flush() + _, err := ai.RunAgentStream(c.Request.Context(), llm, rt, req, func(ev ai.AgentEvent) { + _ = ai.WriteAgentSSE(c.Writer, flusher, ev) + }) + if err != nil && c.Request.Context().Err() == nil { + _ = ai.WriteAgentSSE(c.Writer, flusher, ai.AgentEvent{Type: ai.AgentEventError, Error: err.Error()}) + } + return + } + res, err := ai.RunAgent(c.Request.Context(), llm, rt, req) if err != nil { c.JSON(http.StatusBadGateway, gin.H{"error": err.Error(), "steps": res.Steps, "rounds": res.Rounds}) @@ -499,6 +524,55 @@ func NewRouter(token string) *gin.Engine { c.JSON(http.StatusOK, res) }) + // 用户在气泡确认条点「执行」后：单次写操作（AllowConfirmedWrite），绝不由模型静默触发。 + api.POST("/ai/agent/confirm", func(c *gin.Context) { + var req struct { + ConnectionID string `json:"connectionId"` + SQL string `json:"sql"` + Command string `json:"command"` + Database string `json:"database"` + Kind string `json:"kind"` // query | redis + } + if err := c.ShouldBindJSON(&req); err != nil { + c.JSON(http.StatusBadRequest, gin.H{"error": err.Error()}) + return + } + rt, err := aitools.DefaultRuntime() + if err != nil { + c.JSON(http.StatusInternalServerError, gin.H{"error": err.Error()}) + return + } + if !rt.Policy.WritesEnabled() { + c.JSON(http.StatusForbidden, gin.H{"error": "writes disabled: set mode=read_write in ~/.opsxdb/mcp.json"}) + return + } + rt.Source = "agent" + rt.RejectCallerConfirm = false + rt.AllowConfirmedWrite = true + kind := strings.ToLower(strings.TrimSpace(req.Kind)) + args := map[string]any{"connection_id": req.ConnectionID, "database": req.Database, "confirm_danger": true} + var out any + if kind == "redis" || (kind == "" && req.Command != "") { + args["command"] = req.Command + if args["command"] == "" || args["command"] == nil { + args["command"] = req.SQL + } + out, err = rt.Call(c.Request.Context(), aitools.ToolExecuteRedisCommand, args) + } else { + args["sql"] = req.SQL + out, err = rt.Call(c.Request.Context(), aitools.ToolExecuteQuery, args) + } + if err != nil { + c.JSON(http.StatusBadRequest, gin.H{"error": err.Error()}) + return + } + c.JSON(http.StatusOK, out) + }) + + api.GET("/ai/mcp-status", func(c *gin.Context) { + c.JSON(http.StatusOK, aitools.GetMCPStatus()) + }) + api.POST("/ai/generate", func(c *gin.Context) { var req ai.ChatReq if err := c.ShouldBindJSON(&req); err != nil {