feat: harden send preview idempotency

This commit is contained in:
zhaoyilun
2026-07-10 19:24:53 +08:00
parent 138615a29c
commit 93112d6d99
8 changed files with 482 additions and 22 deletions

View File

@@ -1,6 +1,7 @@
package tools
import (
"bufio"
"context"
"crypto/sha256"
"encoding/hex"
@@ -9,18 +10,21 @@ import (
"os"
"path/filepath"
"strings"
"sync"
"time"
"github.com/modelcontextprotocol/go-sdk/mcp"
)
const (
ToolNameSendMessage = "isphere_send_message"
EnvSendMessageAuditPath = "ISPHERE_SEND_AUDIT_PATH"
sendMessageConnector = "implatform-sidecar"
sendMessagePreviewStage = "preview"
sendMessageBlockedStage = "blocked"
sendMessageProductionMode = "production"
ToolNameSendMessage = "isphere_send_message"
EnvSendMessageAuditPath = "ISPHERE_SEND_AUDIT_PATH"
EnvSendMessageIdempotencyPath = "ISPHERE_SEND_IDEMPOTENCY_PATH"
sendMessageConnector = "implatform-sidecar"
sendMessageConnectorMode = "sandbox-only-shell"
sendMessagePreviewStage = "preview"
sendMessageBlockedStage = "blocked"
sendMessageProductionMode = "production"
)
type SendMessageArgs struct {
@@ -36,6 +40,10 @@ type SendMessageAuditSink interface {
AppendSendMessageAudit(ctx context.Context, event SendMessageAuditEvent) error
}
type SendMessageIdempotencyStore interface {
ReserveSendMessageIdempotency(ctx context.Context, record SendMessageIdempotencyRecord) (SendMessageIdempotencyDecision, error)
}
type SendMessageAuditEvent struct {
Tool string `json:"tool"`
TargetType string `json:"target_type"`
@@ -57,14 +65,46 @@ type SendMessageAuditEvent struct {
IdempotencyKey string `json:"-"`
}
type SendMessageIdempotencyRecord struct {
Tool string `json:"tool"`
TargetRef string `json:"target_ref"`
ExecutionMode string `json:"execution_mode"`
Connector string `json:"connector"`
ConnectorStage string `json:"connector_stage"`
ContentSHA256 string `json:"content_sha256"`
IdempotencyKeySHA256 string `json:"idempotency_key_sha256"`
AuditRef string `json:"audit_ref"`
StartedAt string `json:"started_at"`
}
type SendMessageIdempotencyDecision struct {
Duplicate bool
Conflict bool
FirstAuditRef string
FirstStartedAt string
}
type fileSendMessageAuditSink struct {
path string
}
type fileSendMessageIdempotencyStore struct {
path string
}
var sendMessageIdempotencyFileMu sync.Mutex
func RegisterISphereSendMessageTool(server *mcp.Server, auditSink SendMessageAuditSink) {
RegisterISphereSendMessageToolWithState(server, auditSink, nil)
}
func RegisterISphereSendMessageToolWithState(server *mcp.Server, auditSink SendMessageAuditSink, idempotencyStore SendMessageIdempotencyStore) {
if auditSink == nil {
auditSink = defaultSendMessageAuditSink()
}
if idempotencyStore == nil {
idempotencyStore = defaultSendMessageIdempotencyStore()
}
mcp.AddTool[SendMessageArgs, map[string]any](server, &mcp.Tool{
Name: ToolNameSendMessage,
@@ -77,6 +117,11 @@ func RegisterISphereSendMessageTool(server *mcp.Server, auditSink SendMessageAud
}
finished := time.Now().UTC()
response, event := sendMessagePreviewResponse(normalized, started, finished)
decision, err := idempotencyStore.ReserveSendMessageIdempotency(ctx, sendMessageIdempotencyRecordFromAuditEvent(event))
if err != nil {
return nil, nil, fmt.Errorf("%s idempotency reserve failed: %w", ToolNameSendMessage, err)
}
applySendMessageIdempotencyDecision(response, &event, decision)
if err := auditSink.AppendSendMessageAudit(ctx, event); err != nil {
return nil, nil, fmt.Errorf("%s audit append failed: %w", ToolNameSendMessage, err)
}
@@ -91,6 +136,13 @@ func defaultSendMessageAuditSink() SendMessageAuditSink {
return fileSendMessageAuditSink{path: filepath.Join("runs", "send-audit", "send-message-preview.jsonl")}
}
func defaultSendMessageIdempotencyStore() SendMessageIdempotencyStore {
if path := strings.TrimSpace(os.Getenv(EnvSendMessageIdempotencyPath)); path != "" {
return fileSendMessageIdempotencyStore{path: path}
}
return fileSendMessageIdempotencyStore{path: filepath.Join("runs", "send-audit", "send-message-idempotency.jsonl")}
}
func (s fileSendMessageAuditSink) AppendSendMessageAudit(_ context.Context, event SendMessageAuditEvent) error {
if strings.TrimSpace(s.path) == "" {
return nil
@@ -113,6 +165,83 @@ func (s fileSendMessageAuditSink) AppendSendMessageAudit(_ context.Context, even
return nil
}
func (s fileSendMessageIdempotencyStore) ReserveSendMessageIdempotency(_ context.Context, record SendMessageIdempotencyRecord) (SendMessageIdempotencyDecision, error) {
if strings.TrimSpace(s.path) == "" {
return SendMessageIdempotencyDecision{}, nil
}
if record.IdempotencyKeySHA256 == "" {
return SendMessageIdempotencyDecision{}, fmt.Errorf("idempotency_key_sha256 is required")
}
sendMessageIdempotencyFileMu.Lock()
defer sendMessageIdempotencyFileMu.Unlock()
if err := os.MkdirAll(filepath.Dir(s.path), 0o755); err != nil {
return SendMessageIdempotencyDecision{}, err
}
existing, err := readSendMessageIdempotencyRecords(s.path)
if err != nil {
return SendMessageIdempotencyDecision{}, err
}
for _, candidate := range existing {
if candidate.IdempotencyKeySHA256 != record.IdempotencyKeySHA256 {
continue
}
decision := SendMessageIdempotencyDecision{
FirstAuditRef: candidate.AuditRef,
FirstStartedAt: candidate.StartedAt,
}
if candidate.TargetRef == record.TargetRef && candidate.ContentSHA256 == record.ContentSHA256 {
decision.Duplicate = true
return decision, nil
}
decision.Conflict = true
return decision, nil
}
payload, err := json.Marshal(record)
if err != nil {
return SendMessageIdempotencyDecision{}, err
}
file, err := os.OpenFile(s.path, os.O_CREATE|os.O_APPEND|os.O_WRONLY, 0o600)
if err != nil {
return SendMessageIdempotencyDecision{}, err
}
defer file.Close()
if _, err := file.Write(append(payload, '\n')); err != nil {
return SendMessageIdempotencyDecision{}, err
}
return SendMessageIdempotencyDecision{}, nil
}
func readSendMessageIdempotencyRecords(path string) ([]SendMessageIdempotencyRecord, error) {
file, err := os.Open(path)
if err != nil {
if os.IsNotExist(err) {
return nil, nil
}
return nil, err
}
defer file.Close()
var records []SendMessageIdempotencyRecord
scanner := bufio.NewScanner(file)
for scanner.Scan() {
line := strings.TrimSpace(scanner.Text())
if line == "" {
continue
}
var record SendMessageIdempotencyRecord
if err := json.Unmarshal([]byte(line), &record); err != nil {
return nil, fmt.Errorf("decode idempotency record: %w", err)
}
records = append(records, record)
}
if err := scanner.Err(); err != nil {
return nil, err
}
return records, nil
}
type normalizedSendMessageArgs struct {
TargetType string
TargetID string
@@ -228,12 +357,14 @@ func sendMessagePreviewResponse(input normalizedSendMessageArgs, started time.Ti
SideEffects: sideEffects,
}
response := map[string]any{
"ok": ok,
"send_status": status,
"blocked_reason": blockedReason,
"execution_mode": input.ExecutionMode,
"connector": sendMessageConnector,
"connector_stage": stage,
"ok": ok,
"send_status": status,
"blocked_reason": blockedReason,
"execution_mode": input.ExecutionMode,
"connector": sendMessageConnector,
"connector_mode": sendMessageConnectorMode,
"connector_stage": stage,
"production_send_enabled": false,
"target": map[string]any{
"target_type": input.TargetType,
"target_id": input.TargetID,
@@ -245,6 +376,10 @@ func sendMessagePreviewResponse(input normalizedSendMessageArgs, started time.Ti
},
"idempotency": map[string]any{
"idempotency_key_sha256": input.IdempotencyKeySHA256,
"duplicate_detected": false,
"conflict_detected": false,
"first_audit_ref": nil,
"first_started_at": nil,
},
"side_effects": sideEffects,
"audit": map[string]any{
@@ -252,6 +387,7 @@ func sendMessagePreviewResponse(input normalizedSendMessageArgs, started time.Ti
"source": sendMessageConnector,
"execution_mode": input.ExecutionMode,
"connector": sendMessageConnector,
"connector_mode": sendMessageConnectorMode,
"connector_stage": stage,
"content_sha256": input.ContentSHA256,
"target_ref": input.TargetRef,
@@ -264,6 +400,58 @@ func sendMessagePreviewResponse(input normalizedSendMessageArgs, started time.Ti
return response, event
}
func sendMessageIdempotencyRecordFromAuditEvent(event SendMessageAuditEvent) SendMessageIdempotencyRecord {
return SendMessageIdempotencyRecord{
Tool: event.Tool,
TargetRef: event.TargetRef,
ExecutionMode: event.ExecutionMode,
Connector: event.Connector,
ConnectorStage: event.ConnectorStage,
ContentSHA256: event.ContentSHA256,
IdempotencyKeySHA256: event.IdempotencyKeySHA256,
AuditRef: event.AuditRef,
StartedAt: event.StartedAt,
}
}
func applySendMessageIdempotencyDecision(response map[string]any, event *SendMessageAuditEvent, decision SendMessageIdempotencyDecision) {
idempotency, _ := response["idempotency"].(map[string]any)
if idempotency == nil {
idempotency = map[string]any{}
response["idempotency"] = idempotency
}
idempotency["duplicate_detected"] = decision.Duplicate
idempotency["conflict_detected"] = decision.Conflict
if decision.FirstAuditRef != "" {
idempotency["first_audit_ref"] = decision.FirstAuditRef
}
if decision.FirstStartedAt != "" {
idempotency["first_started_at"] = decision.FirstStartedAt
}
if !decision.Conflict {
return
}
blockedReason := "idempotency key was already used for a different target or content"
response["ok"] = false
response["send_status"] = "blocked"
response["blocked_reason"] = blockedReason
response["connector_stage"] = sendMessageBlockedStage
audit, _ := response["audit"].(map[string]any)
if audit != nil {
audit["result"] = "blocked"
audit["connector_stage"] = sendMessageBlockedStage
audit["blocked_reason"] = blockedReason
}
if event != nil {
event.ConnectorStage = sendMessageBlockedStage
event.SendStatus = "blocked"
event.Result = "blocked"
}
}
func sendMessageNoSideEffects() map[string]any {
return map[string]any{
"sent_message": false,