feat: harden send preview idempotency
This commit is contained in:
@@ -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,
|
||||
|
||||
Reference in New Issue
Block a user