CalvinBackup/CyberStrikeAI
1
0
Fork 0
mirror of https://github.com/Ed1s0nZ/CyberStrikeAI.git synced 2026-09-18 15:42:26 +02:00
Code Issues Packages Projects Releases Wiki Activity

feat: manage task process lifetimes and preserve turn history

Browse Source
This commit is contained in:
Ed1s0nZ
2026-09-16 17:52:58 +08:00
parent fd1c13a43d
commit f7882be546
54 changed files with 3650 additions and 330 deletions
Download Patch File Download Diff File Expand all files Collapse all files
internal/mcp/execution_ownership_test.go
+76
View File
@@ -0,0 +1,76 @@
package mcp
import (
"context"
"cyberstrike-ai/internal/runlease"
"errors"
"testing"
"time"
)
func TestExecutionOwnedAfterContextDetached(t *testing.T) {
scope := runlease.New()
parent, cancel := context.WithCancel(runlease.WithScope(context.Background(), scope))
service := NewExecutionService(nil, nil)
entered := make(chan struct{})
handle, err := service.Submit(parent, ExecutionRequest{Run: func(ctx context.Context) (*ToolResult, error) { close(entered); <-ctx.Done(); return nil, ctx.Err() }})
if err != nil {
t.Fatal(err)
}
<-entered
cancel()
snapshot, _ := service.Get(handle.ID)
if snapshot.Execution.Status != ToolExecutionStatusRunning {
t.Fatal("caller cancellation ended detached worker")
}
scope.Cancel()
deadline, stop := context.WithTimeout(context.Background(), time.Second)
defer stop()
if err = scope.Wait(deadline); err != nil {
t.Fatal(err)
}
snapshot, _ = service.Get(handle.ID)
if snapshot.Execution.Status != ToolExecutionStatusCancelled {
t.Fatalf("unexpected state: %s", snapshot.Execution.Status)
}
if _, err = service.Submit(parent, ExecutionRequest{Run: func(context.Context) (*ToolResult, error) { t.Error("closed task executed tool"); return nil, nil }}); !errors.Is(err, runlease.ErrClosed) {
t.Fatalf("late submit: %v", err)
}
}
func TestRemoteCancellationRequiresAcknowledgement(t *testing.T) {
for _, confirm := range []bool{false, true} {
name := "unconfirmed"
if confirm {
name = "confirmed"
}
t.Run(name, func(t *testing.T) {
scope := runlease.New()
ctx := runlease.WithScope(context.Background(), scope)
service := NewExecutionService(nil, nil)
entered := make(chan struct{})
req := ExecutionRequest{Remote: true, Run: func(ctx context.Context) (*ToolResult, error) { close(entered); <-ctx.Done(); return nil, ctx.Err() }}
if confirm {
req.ConfirmCancellation = func(context.Context) error { return nil }
}
handle, err := service.Submit(ctx, req)
if err != nil {
t.Fatal(err)
}
<-entered
scope.Cancel()
deadline, stop := context.WithTimeout(context.Background(), time.Second)
defer stop()
err = scope.Wait(deadline)
snapshot, _ := service.Get(handle.ID)
if confirm {
if err != nil || snapshot.Execution.Status != ToolExecutionStatusCancelled {
t.Fatalf("confirmed: %v %+v", err, snapshot.Execution)
}
} else {
if !errors.Is(err, runlease.ErrUnconfirmed) || snapshot.Execution.Status != ToolExecutionStatusOrphaned {
t.Fatalf("notification treated as confirmation: %v %+v", err, snapshot.Execution)
}
}
})
}
}
internal/mcp/execution_service.go
+55 -17
View File
@@ -9,6 +9,7 @@ import (
"time"
"cyberstrike-ai/internal/authctx"
"cyberstrike-ai/internal/runlease"
"github.com/google/uuid"
"go.uber.org/zap"
@@ -37,15 +38,18 @@ type ExecutionPreRunFunc func(context.Context, *ToolExecution) (func(), error)
type ExecutionDoneFunc func(*ToolExecution)
type ExecutionRequest struct {
ID string
ToolName string
Arguments map[string]interface{}
ConversationID string
OwnerUserID string
HardTimeout time.Duration
PreRun ExecutionPreRunFunc
Run ExecutionRunFunc
OnDone ExecutionDoneFunc
Remote bool
// A remote adapter may positively confirm server-side cancellation.
ConfirmCancellation func(context.Context) error
ID string
ToolName string
Arguments map[string]interface{}
ConversationID string
OwnerUserID string
HardTimeout time.Duration
PreRun ExecutionPreRunFunc
Run ExecutionRunFunc
OnDone ExecutionDoneFunc
}
type ExecutionHandle struct {
@@ -57,13 +61,17 @@ type ExecutionSnapshot struct {
}
type executionEntry struct {
exec *ToolExecution
cancel context.CancelFunc
done chan struct{}
preRun ExecutionPreRunFunc
run ExecutionRunFunc
result *ToolResult
err error
releaseLease func()
remote bool
runStarted bool
confirmCancellation func(context.Context) error
exec *ToolExecution
cancel context.CancelFunc
done chan struct{}
preRun ExecutionPreRunFunc
run ExecutionRunFunc
result *ToolResult
err error
}
// ExecutionService keeps Eino-facing tool calls synchronous while moving the
@@ -151,12 +159,19 @@ func (s *ExecutionService) Submit(ctx context.Context, req ExecutionRequest) (*E
} else {
runCtx, cancel = context.WithCancel(runCtx)
}
entry := &executionEntry{exec: exec, cancel: cancel, done: make(chan struct{}), preRun: req.PreRun, run: req.Run}
releaseLease, leaseErr := runlease.FromContext(ctx).Register(id, cancel)
if leaseErr != nil {
cancel()
return nil, leaseErr
}
entry := &executionEntry{exec: exec, cancel: cancel, done: make(chan struct{}), preRun: req.PreRun, run: req.Run,
releaseLease: releaseLease, remote: req.Remote, confirmCancellation: req.ConfirmCancellation}
s.mu.Lock()
if _, exists := s.entries[id]; exists {
s.mu.Unlock()
cancel()
releaseLease()
return nil, fmt.Errorf("execution already exists: %s", id)
}
s.entries[id] = entry
@@ -188,8 +203,15 @@ func (s *ExecutionService) runWorker(ctx context.Context, entry *executionEntry,
entry.cancel()
notifyToolRunEnd(ctx, id)
close(entry.done)
if entry.releaseLease != nil {
entry.releaseLease()
}
}()
if ctx.Err() != nil {
s.finishEntry(ctx, entry, nil, ctx.Err(), onDone)
return
}
if entry.preRun != nil {
var preErr error
release, preErr = entry.preRun(ctx, cloneToolExecution(entry.exec))
@@ -198,7 +220,12 @@ func (s *ExecutionService) runWorker(ctx context.Context, entry *executionEntry,
return
}
}
if ctx.Err() != nil {
s.finishEntry(ctx, entry, nil, ctx.Err(), onDone)
return
}
s.markEntryRunning(entry)
entry.runStarted = true
result, err := entryResultRecover(ctx, entry.exec.ToolName, s.logger, func() (*ToolResult, error) {
return nilSafeRun(ctx, entry)
@@ -229,6 +256,12 @@ func (s *ExecutionService) finishEntry(ctx context.Context, entry *executionEntr
if errors.As(err, &blockedErr) {
result, err = blockedErr.result, nil
}
cancellationUnconfirmed := entry.remote && entry.runStarted && ctx.Err() != nil && err != nil
if cancellationUnconfirmed && entry.confirmCancellation != nil {
confirmCtx, cancel := context.WithTimeout(context.Background(), 3*time.Second)
cancellationUnconfirmed = entry.confirmCancellation(confirmCtx) != nil
cancel()
}
cancelledWithUserNote := s.applyAbortUserNoteToCancelledToolResult(id, &result, &err)
now := time.Now()
@@ -287,6 +320,11 @@ func (s *ExecutionService) finishEntry(ctx context.Context, entry *executionEntr
}
entry.exec.Result = result
}
if cancellationUnconfirmed {
entry.exec.Status = ToolExecutionStatusOrphaned
entry.exec.Error = "取消已请求，但远端 MCP 未确认执行已停止"
runlease.FromContext(ctx).MarkUnconfirmed(id, entry.exec.Error)
}
finalExec := cloneToolExecution(entry.exec)
s.mu.Unlock()
internal/mcp/external_manager.go
+14
View File
@@ -706,6 +706,13 @@ func (m *ExternalMCPManager) CallTool(ctx context.Context, toolName string, args
var client ExternalMCPClient
var blockedByGuard bool
handle, err := m.executionService.Submit(ctx, ExecutionRequest{
ConfirmCancellation: func(confirmCtx context.Context) error {
if confirmer, ok := client.(ExternalCancellationConfirmer); ok {
return confirmer.ConfirmToolCancellation(confirmCtx, actualToolName, args)
}
return fmt.Errorf("external MCP client has no cancellation acknowledgement")
},
Remote: true,
ToolName: toolName,
Arguments: args,
ConversationID: MCPConversationIDFromContext(ctx),
@@ -1649,3 +1656,10 @@ func (m *ExternalMCPManager) StopAll() {
}
m.refreshWg.Wait()
}
// ExternalCancellationConfirmer is an optional adapter contract for MCP
// servers with server-side cancellation receipts or lease/task status APIs.
// Ordinary notifications/cancelled must never be treated as confirmation.
type ExternalCancellationConfirmer interface {
ConfirmToolCancellation(context.Context, string, map[string]interface{}) error
}