rlwj-lab/opsxdb
0
0 Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

feat(ai): session harness contract — incomplete, run_id, confirm bind

Browse Source 
Add docs/16 for session state machine, incomplete/resume, queue semantics,
and agent-loop/harness/MCP gap review. Wire RunOutcome on max-rounds,
run_id through audit, TrimMessages/Redact hooks, ConfirmBindHash on write
confirm path; session file skeleton without HTTP yet. No push.
This commit is contained in:
logan committed 2026-10-15 13:58:10 +08:00
1 parent 0aeee2b692
commit 5416e86379
17 files changed
+1041 -61
Download Patch File Download Diff File Expand all files Collapse all files

No files matched your search

server/internal/ai/agent.go
+137 -40
View File
@@ -19,13 +19,17 @@ const (
// AgentReq 侧栏 Agent 请求（与 Ask 的 /api/ai/chat 分路径）。
type AgentReq struct {
Prompt string `json:"prompt"`
ConnectionID string `json:"connectionId,omitempty"`
Driver string `json:"driver,omitempty"`
Database string `json:"database,omitempty"`
Model string `json:"model,omitempty"`
MaxRounds int `json:"maxRounds,omitempty"`
Stream bool `json:"stream,omitempty"`
Prompt string `json:"prompt"`
ConnectionID string `json:"connectionId,omitempty"`
Driver string `json:"driver,omitempty"`
Database string `json:"database,omitempty"`
Model string `json:"model,omitempty"`
MaxRounds int `json:"maxRounds,omitempty"`
Stream bool `json:"stream,omitempty"`
SessionID string `json:"sessionId,omitempty"`
Continue bool `json:"continue,omitempty"` // 真·续跑：注入 Prior
Prior []Message `json:"prior,omitempty"` // 含 tool 轨迹的历史（不含本轮 user）
RunID string `json:"runId,omitempty"` // 可选；空则自动生成
}
// AgentStep 挂在回复下的步骤时间线一项。
@@ -38,16 +42,23 @@ type AgentStep struct {
Ms int64 `json:"ms"`
NeedConfirm bool `json:"needConfirm,omitempty"`
ConfirmSQL string `json:"confirmSql,omitempty"`
ConfirmBind string `json:"confirmBind,omitempty"`
Pending bool `json:"pending,omitempty"`
Result *AgentQueryResult `json:"result,omitempty"`
}
// AgentRes Agent 最终回答 + 步骤。
// AgentRes Agent 最终回答 + 步骤 + 运行结果契约。
type AgentRes struct {
Content string `json:"content"`
SQL string `json:"sql,omitempty"`
Steps []AgentStep `json:"steps"`
Rounds int `json:"rounds"`
Content string `json:"content"`
SQL string `json:"sql,omitempty"`
Steps []AgentStep `json:"steps"`
Rounds int `json:"rounds"`
RunID string `json:"runId,omitempty"`
SessionID string `json:"sessionId,omitempty"`
Outcome string `json:"runOutcome,omitempty"`
IncompleteReason string `json:"incompleteReason,omitempty"`
Resumable bool `json:"resumable,omitempty"`
StepsBudget int `json:"stepsBudget,omitempty"`
}
// RunAgent 多轮：LLM tool_call → aitools.Runtime.Call → 回灌。
@@ -69,39 +80,74 @@ func runAgent(ctx context.Context, llm ToolChat, rt *aitools.Runtime, req AgentR
return AgentRes{}, fmt.Errorf("tool chat LLM required")
}
prompt := strings.TrimSpace(req.Prompt)
if prompt == "" {
if prompt == "" && !(req.Continue && len(req.Prior) > 0) {
return AgentRes{}, fmt.Errorf("prompt required")
}
maxRounds := req.MaxRounds
if maxRounds <= 0 {
maxRounds = defaultAgentMaxRounds
}
if maxRounds > 16 {
maxRounds = 16
}
h := DefaultHarness()
maxRounds := h.ClampRounds(req.MaxRounds)
runID := strings.TrimSpace(req.RunID)
if runID == "" {
runID = NewRunID()
}
rt.RunID = runID
// Agent 路径禁止模型自行带确认令牌
rt.RejectCallerConfirm = true
tools := aitools.OpenAITools()
msgs := []Message{
{Role: "system", Content: buildAgentSystem(req)},
{Role: "user", Content: prompt},
var msgs []Message
if req.Continue && len(req.Prior) > 0 {
msgs = append([]Message{}, req.Prior...)
hasSystem := false
for _, m := range msgs {
if m.Role == "system" {
hasSystem = true
break
}
}
if !hasSystem {
msgs = append([]Message{{Role: "system", Content: buildAgentSystem(req)}}, msgs...)
}
if prompt != "" {
msgs = append(msgs, Message{Role: "user", Content: prompt})
}
msgs = TrimMessages(msgs, h)
} else {
msgs = []Message{
{Role: "system", Content: buildAgentSystem(req)},
{Role: "user", Content: prompt},
}
}
baseMeta := func(steps []AgentStep, rounds int) AgentRes {
return AgentRes{
Steps: steps,
Rounds: rounds,
RunID: runID,
SessionID: req.SessionID,
StepsBudget: maxRounds,
}
}
var steps []AgentStep
var callHistory []string
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})
res := baseMeta(steps, rounds)
res.Content = "已取消"
res.Outcome = OutcomeCancelled
res.Resumable = len(steps) > 0
emitEvent(emit, AgentEvent{Type: AgentEventError, Error: "cancelled", Steps: steps, Rounds: rounds, RunID: runID, Outcome: OutcomeCancelled})
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
res := baseMeta(steps, rounds)
res.Outcome = OutcomeFailed
emitEvent(emit, AgentEvent{Type: AgentEventError, Error: err.Error(), Steps: steps, Rounds: rounds, RunID: runID, Outcome: OutcomeFailed})
return res, err
}
if len(turn.ToolCalls) == 0 {
content := turn.Content
@@ -111,13 +157,14 @@ func runAgent(ctx context.Context, llm ToolChat, rt *aitools.Runtime, req AgentR
for _, part := range chunkDelta(content, 36) {
emitEvent(emit, AgentEvent{Type: AgentEventDelta, Delta: part})
}
res := AgentRes{
Content: content,
SQL: extractSQL(content),
Steps: steps,
Rounds: rounds,
}
emitEvent(emit, AgentEvent{Type: AgentEventDone, Content: res.Content, SQL: res.SQL, Steps: steps, Rounds: rounds})
res := baseMeta(steps, rounds)
res.Content = content
res.SQL = extractSQL(content)
res.Outcome = OutcomeCompleted
emitEvent(emit, AgentEvent{
Type: AgentEventDone, Content: res.Content, SQL: res.SQL, Steps: steps, Rounds: rounds,
RunID: runID, Outcome: OutcomeCompleted, StepsBudget: maxRounds,
})
return res, nil
}
@@ -127,29 +174,76 @@ func runAgent(ctx context.Context, llm ToolChat, rt *aitools.Runtime, req AgentR
ToolCalls: turn.ToolCalls,
})
awaiting := false
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})
argsRaw := strings.TrimSpace(tc.Function.Arguments)
fp := ToolCallFingerprint(name, argsRaw)
if DetectRepeatedToolCall(callHistory, fp, h.SameCallLimit) {
step := AgentStep{
Tool: name,
Args: truncateRunes(argsRaw, maxStepArgsChars),
OK: false,
Summary: "repeated tool call rejected (same args)",
Error: "repeated tool call rejected (same args)",
}
steps = append(steps, step)
cp := step
emitEvent(emit, AgentEvent{Type: AgentEventStep, Step: &cp, RunID: runID})
msgs = append(msgs, Message{
Role: "tool", ToolCallID: tc.ID, Name: name,
Content: step.Error,
})
callHistory = append(callHistory, fp)
continue
}
callHistory = append(callHistory, fp)
pending := AgentStep{Tool: name, Args: truncateRunes(argsRaw, maxStepArgsChars), Pending: true}
emitEvent(emit, AgentEvent{Type: AgentEventStepStart, Step: &pending, RunID: runID})
step, toolMsg := execOneTool(ctx, rt, tc)
steps = append(steps, step)
cp := step
ev := AgentEvent{Type: AgentEventStep, Step: &cp}
ev := AgentEvent{Type: AgentEventStep, Step: &cp, RunID: runID}
if step.Result != nil {
ev.Result = step.Result
}
emitEvent(emit, ev)
msgs = append(msgs, toolMsg)
if step.NeedConfirm {
awaiting = true
}
}
if awaiting {
content := "需要确认写操作后才能继续。"
res := baseMeta(steps, rounds)
res.Content = content
res.Outcome = OutcomeAwaitingConfirm
res.Resumable = true
emitEvent(emit, AgentEvent{
Type: AgentEventDone, Content: content, Steps: steps, Rounds: rounds,
RunID: runID, Outcome: OutcomeAwaitingConfirm, Resumable: true, StepsBudget: maxRounds,
})
return res, nil
}
msgs = TrimMessages(msgs, h)
}
content := fmt.Sprintf("已达最大工具轮次（%d）。请根据已执行步骤自行判断，或缩小问题后重试。", maxRounds)
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})
res := baseMeta(steps, rounds)
res.Content = content
res.Outcome = OutcomeIncomplete
res.IncompleteReason = ReasonMaxSteps
res.Resumable = true
emitEvent(emit, AgentEvent{
Type: AgentEventIncomplete, Content: content, Steps: steps, Rounds: rounds,
RunID: runID, Outcome: OutcomeIncomplete, IncompleteReason: ReasonMaxSteps,
Resumable: true, StepsBudget: maxRounds,
})
return res, nil
}
@@ -228,6 +322,9 @@ func execOneTool(ctx context.Context, rt *aitools.Runtime, tc ToolCall) (AgentSt
} else if cmd, _ := m["command"].(string); cmd != "" {
step.ConfirmSQL = cmd
}
if bind, _ := m["confirmBind"].(string); bind != "" {
step.ConfirmBind = bind
}
msg, _ := m["message"].(string)
if msg == "" {
msg = "requires user confirmation (confirm_danger)"
server/internal/ai/agent_stream.go
+16 -10
View File
@@ -12,7 +12,8 @@ const (
AgentEventStep = "step"
AgentEventDelta = "delta"
AgentEventDone = "done"
AgentEventError = "error"
AgentEventError = "error"
AgentEventIncomplete = "incomplete"
)
// AgentQueryResult 可供「在结果区打开」的行列结果。
@@ -26,15 +27,20 @@ type AgentQueryResult struct {
// 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"`
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"`
RunID string `json:"runId,omitempty"`
Outcome string `json:"runOutcome,omitempty"`
IncompleteReason string `json:"incompleteReason,omitempty"`
Resumable bool `json:"resumable,omitempty"`
StepsBudget int `json:"stepsBudget,omitempty"`
}
// AgentEmitter 可选；nil 时仅返回最终 AgentRes（非流式 fallback）。
server/internal/ai/harness.go
+165
View File
@@ -0,0 +1,165 @@
package ai
import (
"crypto/rand"
"encoding/hex"
"encoding/json"
"strings"
)
// HarnessConfig Agent 运行时预算（桌面裁剪版；不做 Evolver）。
type HarnessConfig struct {
MaxRounds int // 默认 8，硬帽 16
MaxToolResultChars int // 回灌截断
MaxContextTokens int // trim 预算（粗估字符/4）
SameCallLimit int // 同工具同参上限；0=关闭
MinRetainMessages int // trim 至少保留条数（不含 system）
}
// DefaultHarness 与历史 agent.go 常量对齐。
func DefaultHarness() HarnessConfig {
return HarnessConfig{
MaxRounds: defaultAgentMaxRounds,
MaxToolResultChars: maxToolResultFeedChars,
MaxContextTokens: 8000,
SameCallLimit: 3,
MinRetainMessages: 4,
}
}
func (h HarnessConfig) ClampRounds(n int) int {
if n <= 0 {
n = h.MaxRounds
if n <= 0 {
n = defaultAgentMaxRounds
}
}
if n > 16 {
n = 16
}
return n
}
// NewRunID 生成短 run_id，贯穿审计与步骤。
func NewRunID() string {
var b [8]byte
_, _ = rand.Read(b[:])
return "run_" + hex.EncodeToString(b[:])
}
// EstimateTokens 粗估（字符/4）；含 tool_calls JSON。
func EstimateTokens(msgs []Message) int {
n := 0
for _, m := range msgs {
n += len(m.Content)
if len(m.ToolCalls) > 0 {
if b, err := json.Marshal(m.ToolCalls); err == nil {
n += len(b)
}
}
n += len(m.Name) + len(m.ToolCallID) + 8
}
if n < 0 {
return 0
}
return n / 4
}
// TrimMessages 保留 system + 最近消息，控制在 MaxContextTokens*90%；
// 并按 tool_call_id 整组保留 assistant.tool_calls 与对应 tool。
func TrimMessages(msgs []Message, h HarnessConfig) []Message {
if len(msgs) == 0 {
return msgs
}
budget := h.MaxContextTokens
if budget <= 0 {
budget = 8000
}
limit := budget * 90 / 100
minRetain := h.MinRetainMessages
if minRetain <= 0 {
minRetain = 4
}
var system []Message
var rest []Message
for _, m := range msgs {
if m.Role == "system" {
system = append(system, m)
} else {
rest = append(rest, m)
}
}
if EstimateTokens(msgs) <= limit {
return ensureToolPairing(msgs)
}
// 从尾部向前取，直到接近预算
var kept []Message
for i := len(rest) - 1; i >= 0; i-- {
candidate := append([]Message{}, rest[i:]...)
candidate = append(append([]Message{}, system...), candidate...)
if EstimateTokens(candidate) > limit && len(kept) >= minRetain {
break
}
kept = rest[i:]
if len(kept) >= minRetain && EstimateTokens(candidate) <= limit {
break
}
}
if len(kept) == 0 && len(rest) > 0 {
// 至少留最后一条
kept = rest[len(rest)-1:]
}
out := append(append([]Message{}, system...), kept...)
return ensureToolPairing(out)
}
// ensureToolPairing 丢弃没有对应 assistant tool_calls 的孤立 tool，以及裁切后半截的 tool_calls。
func ensureToolPairing(msgs []Message) []Message {
toolIDs := map[string]bool{}
for _, m := range msgs {
if m.Role == "assistant" {
for _, tc := range m.ToolCalls {
if tc.ID != "" {
toolIDs[tc.ID] = true
}
}
}
}
out := make([]Message, 0, len(msgs))
for _, m := range msgs {
if m.Role == "tool" {
if m.ToolCallID == "" || !toolIDs[m.ToolCallID] {
continue
}
}
if m.Role == "assistant" && len(m.ToolCalls) > 0 {
// 若后续缺少任一 tool 结果，整组 assistant 仍保留（由调用方补 incomplete）；
// 这里只滤 tool 侧孤儿。
}
out = append(out, m)
}
return out
}
// ToolCallFingerprint 同参循环检测键。
func ToolCallFingerprint(name, argsJSON string) string {
name = strings.TrimSpace(name)
argsJSON = strings.Join(strings.Fields(argsJSON), " ")
return name + "\n" + argsJSON
}
// DetectRepeatedToolCall 若同一指纹已出现 limit 次（不含本次），返回 true。
func DetectRepeatedToolCall(history []string, fp string, limit int) bool {
if limit <= 0 || fp == "" {
return false
}
n := 0
for _, h := range history {
if h == fp {
n++
}
}
return n >= limit
}
server/internal/ai/harness_test.go
+155
View File
@@ -0,0 +1,155 @@
package ai
import (
"context"
"encoding/json"
"strings"
"testing"
"opsxdb/server/internal/aitools"
)
func TestTrimMessagesKeepsSystemAndDropsOrphanTool(t *testing.T) {
h := DefaultHarness()
h.MaxContextTokens = 40
msgs := []Message{
{Role: "system", Content: "sys"},
{Role: "user", Content: strings.Repeat("u", 400)},
{Role: "assistant", Content: "a", ToolCalls: []ToolCall{{ID: "c1", Type: "function", Function: ToolCallFunc{Name: "opsxdb_list_tables", Arguments: "{}"}}}},
{Role: "tool", ToolCallID: "c1", Name: "opsxdb_list_tables", Content: `{"tables":["users"]}`},
{Role: "user", Content: "continue please"},
}
out := TrimMessages(msgs, h)
hasSys := false
for _, m := range out {
if m.Role == "system" {
hasSys = true
}
}
if !hasSys {
t.Fatal("system dropped")
}
orphan := TrimMessages([]Message{
{Role: "system", Content: "s"},
{Role: "tool", ToolCallID: "x", Content: "nope"},
{Role: "user", Content: "hi"},
}, h)
for _, m := range orphan {
if m.Role == "tool" {
t.Fatalf("expected orphan tool dropped: %+v", orphan)
}
}
}
func TestDetectRepeatedToolCall(t *testing.T) {
fp := ToolCallFingerprint("opsxdb_list_tables", `{"connection_id":"c1"}`)
hist := []string{fp, fp, fp}
if !DetectRepeatedToolCall(hist, fp, 3) {
t.Fatal("want reject at 3")
}
if DetectRepeatedToolCall(hist[:2], fp, 3) {
t.Fatal("should allow under limit")
}
}
func TestRedactSensitive(t *testing.T) {
out := RedactSensitive("Bearer tokensecret password=secret mysql://user:pass@host/db")
if strings.Contains(out, "tokensecret") || strings.Contains(out, "password=secret") || strings.Contains(out, ":pass@") {
t.Fatalf("not redacted: %s", out)
}
if !strings.Contains(RedactSensitive("sk-abcdefghijklmnop"), "[REDACTED]") {
t.Fatal("sk key")
}
}
func TestRunAgentMaxRoundsIncomplete(t *testing.T) {
rt := aitools.NewRuntime(aitools.Policy{
Mode: aitools.ModeReadOnly,
AllowedConnectionIDs: []string{"c1"},
DefaultLimit: 100,
MaxLimit: 500,
})
rt.LoadConns = func() ([]aitools.SavedConnection, error) {
return []aitools.SavedConnection{{ID: "c1", Name: "d", Driver: "mysql"}}, nil
}
sawRunID := false
rt.Audit = func(e aitools.AuditEntry) {
if e.RunID != "" {
sawRunID = true
}
}
var turns []ToolsTurn
for i := 0; i < 3; i++ {
turns = append(turns, ToolsTurn{ToolCalls: []ToolCall{{
ID: "c", Type: "function",
Function: ToolCallFunc{Name: aitools.ToolListConnections, Arguments: "{}"},
}}})
}
llm := &scriptedLLM{turns: turns}
res, err := RunAgent(context.Background(), llm, rt, AgentReq{Prompt: "loop", MaxRounds: 2})
if err != nil {
t.Fatal(err)
}
if res.Outcome != OutcomeIncomplete {
t.Fatalf("outcome=%s want incomplete", res.Outcome)
}
if res.IncompleteReason != ReasonMaxSteps {
t.Fatalf("reason=%s", res.IncompleteReason)
}
if !res.Resumable || res.RunID == "" {
t.Fatalf("resumable/runId: %+v", res)
}
if !sawRunID {
t.Fatal("expected audit run_id")
}
}
func TestRunAgentCompletedHasOutcome(t *testing.T) {
rt := aitools.NewRuntime(aitools.Policy{Mode: aitools.ModeReadOnly, AllowedConnectionIDs: []string{"c1"}, DefaultLimit: 100, MaxLimit: 500})
rt.Audit = func(aitools.AuditEntry) {}
llm := &scriptedLLM{turns: []ToolsTurn{{Content: "hello"}}}
res, err := RunAgent(context.Background(), llm, rt, AgentReq{Prompt: "hi"})
if err != nil {
t.Fatal(err)
}
if res.Outcome != OutcomeCompleted {
t.Fatalf("outcome=%s", res.Outcome)
}
if res.RunID == "" || res.StepsBudget != defaultAgentMaxRounds {
t.Fatalf("%+v", res)
}
}
func TestRunAgentRepeatCallRejected(t *testing.T) {
rt := aitools.NewRuntime(aitools.Policy{Mode: aitools.ModeReadOnly, AllowedConnectionIDs: []string{"c1"}, DefaultLimit: 100, MaxLimit: 500})
rt.LoadConns = func() ([]aitools.SavedConnection, error) {
return []aitools.SavedConnection{{ID: "c1", Name: "d", Driver: "mysql"}}, nil
}
rt.Audit = func(aitools.AuditEntry) {}
args := `{"connection_id":"c1"}`
var turns []ToolsTurn
// 同一轮多个相同 tool_call，以及跨轮：SameCallLimit=3 指 history 已有 3 次则拒第 4 次
// 每轮一个相同调用，第 4 轮应被拒但仍继续到最终答案
for i := 0; i < 4; i++ {
turns = append(turns, ToolsTurn{ToolCalls: []ToolCall{{
ID: "c" + string(rune('1'+i)), Type: "function",
Function: ToolCallFunc{Name: aitools.ToolListConnections, Arguments: args},
}}})
}
turns = append(turns, ToolsTurn{Content: "done"})
llm := &scriptedLLM{turns: turns}
res, err := RunAgent(context.Background(), llm, rt, AgentReq{Prompt: "x", MaxRounds: 8})
if err != nil {
t.Fatal(err)
}
rejected := 0
for _, s := range res.Steps {
if strings.Contains(s.Summary, "repeated tool call") {
rejected++
}
}
if rejected < 1 {
t.Fatalf("expected at least one repeated rejection, steps=%+v", res.Steps)
}
_ = json.Marshal
}
server/internal/ai/outcome.go
+19
View File
@@ -0,0 +1,19 @@
package ai
// RunOutcome 单次 Agent 运行结束状态（对齐 aiops docs/ai-agent/13）。
const (
OutcomeCompleted = "completed"
OutcomeIncomplete = "incomplete"
OutcomeFailed = "failed"
OutcomeCancelled = "cancelled"
OutcomeAwaitingConfirm = "awaiting_confirm"
)
// IncompleteReason 仅当 Outcome==incomplete 时有意义。
const (
ReasonMaxSteps = "max_steps"
ReasonContextBudget = "context_budget"
ReasonToolTimeout = "tool_timeout"
ReasonLLMTimeout = "llm_timeout"
ReasonSafetyAbort = "safety_abort"
)
server/internal/ai/redact.go
+42
View File
@@ -0,0 +1,42 @@
package ai
import (
"regexp"
"strings"
)
var (
reBearer = regexp.MustCompile(`(?i)(Bearer\s+)[A-Za-z0-9._\-]+`)
reSkKey = regexp.MustCompile(`(?i)\b(sk-[A-Za-z0-9]{8,})`)
rePassword = regexp.MustCompile(`(?i)(password|passwd|pwd)\s*[=:]\s*\S+`)
reDSN = regexp.MustCompile(`(?i)\b((?:mysql|postgres|postgresql|mongodb|redis)://[^:\s/]+):([^@\s]+)@`)
rePEM = regexp.MustCompile(`(?s)-----BEGIN [A-Z0-9 ]*PRIVATE KEY-----.*?-----END [A-Z0-9 ]*PRIVATE KEY-----`)
)
// RedactSensitive 脱敏：落盘 / 审计 summary / 展示共用。
func RedactSensitive(s string) string {
if s == "" {
return s
}
s = rePEM.ReplaceAllString(s, "[REDACTED_PRIVATE_KEY]")
s = reDSN.ReplaceAllString(s, "${1}:[REDACTED]@")
s = reBearer.ReplaceAllString(s, "${1}[REDACTED]")
s = reSkKey.ReplaceAllString(s, "sk-[REDACTED]")
s = rePassword.ReplaceAllString(s, "${1}=[REDACTED]")
return s
}
// RedactMessages 对消息 Content 做脱敏（不改 tool_call 结构）。
func RedactMessages(msgs []Message) []Message {
out := make([]Message, len(msgs))
copy(out, msgs)
for i := range out {
out[i].Content = RedactSensitive(out[i].Content)
}
return out
}
// NormalizeStatement 确认绑定前规范化 SQL/命令。
func NormalizeStatement(s string) string {
return strings.Join(strings.Fields(strings.TrimSpace(s)), " ")
}
server/internal/ai/session.go
+89
View File
@@ -0,0 +1,89 @@
package ai
import (
"encoding/json"
"fmt"
"os"
"path/filepath"
"time"
)
// Session 桌面本地会话（目标态落盘；HTTP 路由可后挂）。
type Session struct {
ID string `json:"id"`
Title string `json:"title,omitempty"`
Status string `json:"status"` // idle|running|incomplete|completed|failed|cancelled
CreatedAt string `json:"createdAt"`
UpdatedAt string `json:"updatedAt"`
Messages []Message `json:"messages"`
Runs []SessionRun `json:"runs,omitempty"`
}
// SessionRun 一次运行摘要。
type SessionRun struct {
RunID string `json:"runId"`
Outcome string `json:"outcome"`
IncompleteReason string `json:"incompleteReason,omitempty"`
Resumable bool `json:"resumable"`
StepsUsed int `json:"stepsUsed"`
StepsBudget int `json:"stepsBudget"`
FinishedAt string `json:"finishedAt"`
}
func sessionsDir() (string, error) {
home, err := os.UserHomeDir()
if err != nil {
return "", err
}
return filepath.Join(home, ".opsxdb", "ai-sessions"), nil
}
// SessionPath 返回会话文件路径。
func SessionPath(id string) (string, error) {
dir, err := sessionsDir()
if err != nil {
return "", err
}
return filepath.Join(dir, id+".json"), nil
}
// SaveSession 脱敏后写入磁盘。
func SaveSession(s Session) error {
if s.ID == "" {
return fmt.Errorf("session id required")
}
s.Messages = RedactMessages(s.Messages)
s.UpdatedAt = time.Now().Format(time.RFC3339)
if s.CreatedAt == "" {
s.CreatedAt = s.UpdatedAt
}
path, err := SessionPath(s.ID)
if err != nil {
return err
}
if err := os.MkdirAll(filepath.Dir(path), 0700); err != nil {
return err
}
b, err := json.MarshalIndent(s, "", " ")
if err != nil {
return err
}
return os.WriteFile(path, b, 0600)
}
// LoadSession 读会话；不存在返回 os.ErrNotExist。
func LoadSession(id string) (Session, error) {
path, err := SessionPath(id)
if err != nil {
return Session{}, err
}
b, err := os.ReadFile(path)
if err != nil {
return Session{}, err
}
var s Session
if err := json.Unmarshal(b, &s); err != nil {
return Session{}, err
}
return s, nil
}
server/internal/aitools/audit.go
+1
View File
@@ -15,6 +15,7 @@ type AuditEntry struct {
ConnectionID string `json:"connection_id,omitempty"`
Mode string `json:"mode"`
Source string `json:"source,omitempty"` // mcp | agent
RunID string `json:"run_id,omitempty"`
OK bool `json:"ok"`
Summary string `json:"summary,omitempty"`
Error string `json:"error,omitempty"`
server/internal/aitools/confirm_bind.go
+29
View File
@@ -0,0 +1,29 @@
package aitools
import (
"crypto/sha256"
"encoding/hex"
"fmt"
"strings"
)
// ConfirmBindHash 绑定 connection + 规范化语句；改语句则哈希变，须重确认。
func ConfirmBindHash(connectionID, statement string) string {
connectionID = strings.TrimSpace(connectionID)
statement = strings.Join(strings.Fields(strings.TrimSpace(statement)), " ")
sum := sha256.Sum256([]byte(connectionID + "\n" + statement))
return hex.EncodeToString(sum[:])
}
// VerifyConfirmBind 校验客户端带回的 bind；空 bind 视为不合法（桌面路径必须带）。
func VerifyConfirmBind(connectionID, statement, bind string) error {
bind = strings.TrimSpace(bind)
if bind == "" {
return fmt.Errorf("confirmBind required")
}
want := ConfirmBindHash(connectionID, statement)
if !strings.EqualFold(bind, want) {
return fmt.Errorf("confirmBind mismatch: statement or connection changed; re-confirm required")
}
return nil
}
server/internal/aitools/confirm_bind_test.go
+57
View File
@@ -0,0 +1,57 @@
package aitools
import "testing"
func TestConfirmBindHashStableAndSensitive(t *testing.T) {
a := ConfirmBindHash("c1", "INSERT INTO t VALUES (1)")
b := ConfirmBindHash("c1", " INSERT INTO t VALUES (1) ")
if a != b {
t.Fatalf("normalize mismatch %s vs %s", a, b)
}
c := ConfirmBindHash("c1", "INSERT INTO t VALUES (2)")
if a == c {
t.Fatal("different sql should differ")
}
if err := VerifyConfirmBind("c1", "INSERT INTO t VALUES (1)", a); err != nil {
t.Fatal(err)
}
if err := VerifyConfirmBind("c1", "INSERT INTO t VALUES (2)", a); err == nil {
t.Fatal("want mismatch")
}
if err := VerifyConfirmBind("c1", "INSERT INTO t VALUES (1)", ""); err == nil {
t.Fatal("want empty bind error")
}
}
func TestNeedConfirmIncludesBind(t *testing.T) {
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: "127.0.0.1"}, nil
}
rt.Audit = func(AuditEntry) {}
out, err := rt.Call(nil, ToolExecuteQuery, map[string]any{
"connection_id": "c1",
"sql": "DELETE FROM users WHERE id=1",
})
if err != nil {
t.Fatal(err)
}
m, ok := out.(map[string]any)
if !ok || m["needConfirm"] != true {
t.Fatalf("%v", out)
}
bind, _ := m["confirmBind"].(string)
if bind == "" {
t.Fatal("missing confirmBind")
}
want := ConfirmBindHash("c1", "DELETE FROM users WHERE id=1")
if bind != want {
t.Fatalf("bind %s want %s", bind, want)
}
}
server/internal/aitools/runtime.go
+17 -8
View File
@@ -53,6 +53,8 @@ type Runtime struct {
AllowConfirmedWrite bool
// Source 写入审计：mcp | agent。
Source string
// RunID 写入审计，串联一次 Agent run。
RunID string
}
func NewRuntime(policy Policy) *Runtime {
@@ -77,6 +79,9 @@ func (r *Runtime) audit(e AuditEntry) {
if e.Source == "" {
e.Source = r.Source
}
if e.RunID == "" && r.RunID != "" {
e.RunID = r.RunID
}
if r.Audit != nil {
r.Audit(e)
return
@@ -399,11 +404,13 @@ func (r *Runtime) executeQuery(ctx context.Context, args map[string]any) (any, e
// P1：可写已开启，但无确认令牌则返回 needConfirm（永不静默写）。
if !confirmed {
r.audit(AuditEntry{Tool: ToolExecuteQuery, ConnectionID: id, OK: false, Error: "needConfirm", Summary: truncate(sqlText, 80)})
bind := ConfirmBindHash(id, sqlText)
return map[string]any{
"needConfirm": true,
"sql": sqlText,
"message": "write requires explicit user confirmation",
"reason": readReason,
"needConfirm": true,
"sql": sqlText,
"confirmBind": bind,
"message": "write requires explicit user confirmation",
"reason": readReason,
}, nil
}
}
@@ -471,11 +478,13 @@ func (r *Runtime) executeRedis(ctx context.Context, args map[string]any) (any, e
}
if !confirmed {
r.audit(AuditEntry{Tool: ToolExecuteRedisCommand, ConnectionID: id, OK: false, Error: "needConfirm", Summary: truncate(command, 80)})
bind := ConfirmBindHash(id, command)
return map[string]any{
"needConfirm": true,
"command": command,
"message": "redis write requires explicit user confirmation",
"reason": rerr.Error(),
"needConfirm": true,
"command": command,
"confirmBind": bind,
"message": "redis write requires explicit user confirmation",
"reason": rerr.Error(),
}, nil
}
}
server/internal/server/router.go
+11
View File
@@ -532,6 +532,8 @@ func NewRouter(token string) *gin.Engine {
Command string `json:"command"`
Database string `json:"database"`
Kind string `json:"kind"` // query | redis
ConfirmBind string `json:"confirmBind"`
RunID string `json:"runId"`
}
if err := c.ShouldBindJSON(&req); err != nil {
c.JSON(http.StatusBadRequest, gin.H{"error": err.Error()})
@@ -546,7 +548,16 @@ func NewRouter(token string) *gin.Engine {
c.JSON(http.StatusForbidden, gin.H{"error": "writes disabled: set mode=read_write in ~/.opsxdb/mcp.json"})
return
}
stmt := strings.TrimSpace(req.SQL)
if stmt == "" {
stmt = strings.TrimSpace(req.Command)
}
if err := aitools.VerifyConfirmBind(req.ConnectionID, stmt, req.ConfirmBind); err != nil {
c.JSON(http.StatusBadRequest, gin.H{"error": err.Error()})
return
}
rt.Source = "agent"
rt.RunID = strings.TrimSpace(req.RunID)
rt.RejectCallerConfirm = false
rt.AllowConfirmedWrite = true
kind := strings.ToLower(strings.TrimSpace(req.Kind))