From 19b60132da8c8d31577b46b74471c38bbd6172ff Mon Sep 17 00:00:00 2001 From: Hz_ Date: Tue, 28 Jul 2026 14:59:34 +0800 Subject: [PATCH] fix(go-agent): use session IDs for cancellation and context flow (#17462) ## Summary - Propagate request contexts through Agent Canvas execution and external calls. - Replace internal task IDs with session IDs while retaining `task_id` as a wire alias. - Complete session-scoped cancellation with Redis lease and token validation. ## Testing - Go backend tests passed. image --- internal/agent/canvas/cancel.go | 91 +++-- internal/agent/canvas/cancel_test.go | 41 ++- internal/agent/canvas/canvas.go | 4 +- internal/agent/canvas/compile.go | 8 +- internal/agent/canvas/parallel_subgraph.go | 4 +- internal/agent/canvas/run_tracker.go | 182 +++++++++- internal/agent/canvas/run_tracker_test.go | 155 +++++++++ internal/agent/canvas/runner.go | 113 ++---- internal/agent/canvas/runner_cancel_test.go | 68 ++++ internal/agent/canvas/scheduler.go | 19 +- .../agent/canvas/state_serializer_test.go | 14 +- internal/agent/canvas/stream.go | 6 +- internal/agent/canvas/stream_test.go | 16 +- internal/agent/component/agent.go | 19 +- internal/agent/component/agent_test.go | 8 +- internal/agent/component/message.go | 4 +- .../agent/component/message_phase8b_test.go | 11 +- internal/agent/component/stagehand_runtime.go | 8 +- .../agent/component/tool_dispatch_test.go | 4 +- internal/agent/dsl/reset.go | 6 +- internal/agent/runtime/state.go | 12 +- internal/agent/runtime/state_test.go | 8 +- internal/agent/sandbox/manager_client.go | 7 +- internal/agent/sandbox/ssh.go | 99 ++++-- internal/dao/api_token.go | 14 + internal/handler/agent.go | 38 +- internal/handler/agent_test.go | 60 +++- internal/handler/agent_wait_for_user_test.go | 1 - internal/handler/agent_webhook.go | 52 ++- internal/router/agent_routes.go | 11 +- internal/router/agent_routes_test.go | 20 +- internal/router/router.go | 3 + internal/service/agent.go | 326 +++++++++++++++--- internal/service/agent_cancel_test.go | 174 ++++++++++ internal/service/agent_test.go | 2 +- internal/service/bot_completion.go | 18 +- .../service/bot_completion_history_test.go | 21 +- 37 files changed, 1323 insertions(+), 324 deletions(-) create mode 100644 internal/agent/canvas/runner_cancel_test.go create mode 100644 internal/service/agent_cancel_test.go diff --git a/internal/agent/canvas/cancel.go b/internal/agent/canvas/cancel.go index 76ed925712..6afcda9626 100644 --- a/internal/agent/canvas/cancel.go +++ b/internal/agent/canvas/cancel.go @@ -14,26 +14,29 @@ // limitations under the License. // -// cancel.go implements the cross-process cancel signal. See plan §4.9 — -// a Go canvas run goroutine polls Redis for "{taskID}-cancel"; when the -// HTTP handler sets the key, the watcher fires onCancel. The Redis key -// naming is deliberately identical to the Python task_service.py -// protocol (line 521-523) so Go and Python canvas runs in the same -// tenant can signal each other. +// cancel.go implements the cross-process cancel signal for ordinary Agent +// sessions. A canvas run polls Redis for "{sessionID}-cancel"; when the HTTP +// handler sets the key, the watcher cancels the same context used by the +// workflow. package canvas import ( "context" "errors" - redis2 "ragflow/internal/engine/redis" "time" "github.com/redis/go-redis/v9" + "go.uber.org/zap" + + "ragflow/internal/common" + redis2 "ragflow/internal/engine/redis" ) -// cancelKeySuffix is appended to the task id to form the Redis key. +// cancelKeySuffix is appended to the session id to form the Redis key. const cancelKeySuffix = "-cancel" +func cancelKey(sessionID string) string { return sessionID + cancelKeySuffix } + // cancelPollInterval is the gap between Redis Get polls. 500ms keeps // cancel latency p99 ≤ 500ms while staying cheap (one GET every half- // second per active run). Tunable later if a tenant needs lower latency. @@ -60,30 +63,24 @@ var cancelClientFn = func() (*redis.Client, error) { } // WatchCancel blocks until either ctx is cancelled or the Redis -// "{taskID}-cancel" key is set to a non-empty value. When fired, it -// calls onCancel exactly once and returns. Polling interval is fixed +// "{sessionID}-cancel" key contains a non-empty value. When fired, it calls +// onCancel exactly once and returns. Polling interval is fixed // at 500ms (see plan §4.9 — revised 2026-06-03 from 1s to 500ms). // -// WatchCancel is intended to run as a side goroutine; the run-loop -// goroutine calls it with onCancel wired to the eino graph interrupt -// callback: +// WatchCancel is intended to run as a side goroutine with onCancel wired to +// the cancel function for the workflow context: // -// go func() { -// canvas.WatchCancel(ctx, taskID, func() { -// interrupt(compose.WithGraphInterruptTimeout(30*time.Second)) -// }) -// }() -func WatchCancel(ctx context.Context, taskID string, onCancel func()) { +// go canvas.WatchCancel(ctx, sessionID, cancelRun) +func WatchCancel(ctx context.Context, sessionID string, onCancel func()) { c, err := cancelClientFn() if err != nil { - // Without Redis the watcher can do nothing. Returning silently - // matches the rest of the canvas layer: a missing cache is a - // deployment error surfaced at startup, not at every call. + common.Warn("agent cancel watcher unavailable", zap.String("session_id", sessionID), zap.Error(err)) return } - key := taskID + cancelKeySuffix + key := cancelKey(sessionID) ticker := time.NewTicker(cancelPollInterval) defer ticker.Stop() + readFailureLogged := false for { select { @@ -92,11 +89,13 @@ func WatchCancel(ctx context.Context, taskID string, onCancel func()) { case <-ticker.C: v, err := c.Get(ctx, key).Result() if err != nil && !errors.Is(err, redis.Nil) { - // Transient Redis error — log by skipping this tick; the - // next tick will retry. Avoid spinning on persistent - // failure. + if !readFailureLogged { + common.Warn("agent cancel watcher redis read failed", zap.String("session_id", sessionID), zap.Error(err)) + readFailureLogged = true + } continue } + readFailureLogged = false if v != "" { if onCancel != nil { onCancel() @@ -107,14 +106,42 @@ func WatchCancel(ctx context.Context, taskID string, onCancel func()) { } } -// RequestCancel publishes a cancel signal for the given task. The -// 24h TTL matches the Python task_service.py protocol so a flag set -// during one run is still observable by a resume that arrives hours -// later (e.g. after a long client-side wait). -func RequestCancel(ctx context.Context, taskID string) error { +// RequestCancel publishes a cancel signal for the given session. The marker +// is cleared atomically when the owning run releases its active lease. +func RequestCancel(ctx context.Context, sessionID string) error { c, err := cancelClientFn() if err != nil { return err } - return c.Set(ctx, taskID+cancelKeySuffix, "x", RequestCancelTTL).Err() + return c.Set(ctx, cancelKey(sessionID), "x", RequestCancelTTL).Err() +} + +// CancelRequested reports whether a cancellation marker is currently +// published for sessionID. Registration uses this as an immediate check after +// acquiring the lease so a cancel that wins the startup race cannot be erased +// before the polling watcher gets its first tick. +func CancelRequested(ctx context.Context, sessionID string) (bool, error) { + c, err := cancelClientFn() + if err != nil { + return false, err + } + v, err := c.Get(ctx, cancelKey(sessionID)).Result() + if errors.Is(err, redis.Nil) { + return false, nil + } + if err != nil { + return false, err + } + return v != "", nil +} + +// ClearCancel removes a session cancel marker. Production run registration +// and release perform this operation atomically with active-session ownership; +// this helper remains useful for Redis-degraded and focused test paths. +func ClearCancel(ctx context.Context, sessionID string) error { + c, err := cancelClientFn() + if err != nil { + return err + } + return c.Del(ctx, cancelKey(sessionID)).Err() } diff --git a/internal/agent/canvas/cancel_test.go b/internal/agent/canvas/cancel_test.go index 13eedc3304..713d10c8b7 100644 --- a/internal/agent/canvas/cancel_test.go +++ b/internal/agent/canvas/cancel_test.go @@ -49,18 +49,18 @@ func TestWatchCancel_FiresAfterRequest(t *testing.T) { withCancelClient(t) ctx := t.Context() - taskID := "task_test_1" + sessionID := "session_test_1" fired := atomic.Bool{} done := make(chan struct{}) go func() { - WatchCancel(ctx, taskID, func() { fired.Store(true) }) + WatchCancel(ctx, sessionID, func() { fired.Store(true) }) close(done) }() // Give the watcher time to start its first tick. time.Sleep(200 * time.Millisecond) - if err := RequestCancel(ctx, taskID); err != nil { + if err := RequestCancel(ctx, sessionID); err != nil { t.Fatalf("RequestCancel: %v", err) } @@ -80,10 +80,10 @@ func TestWatchCancel_StopsOnContextCancel(t *testing.T) { withCancelClient(t) ctx, cancel := context.WithCancel(context.Background()) - taskID := "task_test_ctx" + sessionID := "session_test_ctx" done := make(chan struct{}) go func() { - WatchCancel(ctx, taskID, func() { + WatchCancel(ctx, sessionID, func() { t.Error("onCancel should not fire without a Redis signal") }) close(done) @@ -109,7 +109,7 @@ func TestWatchCancel_OnCancelNotInvokedForEmptyKey(t *testing.T) { invoked := atomic.Int32{} done := make(chan struct{}) go func() { - WatchCancel(ctx, "task_never_cancelled", func() { + WatchCancel(ctx, "session_never_cancelled", func() { invoked.Add(1) }) close(done) @@ -121,28 +121,47 @@ func TestWatchCancel_OnCancelNotInvokedForEmptyKey(t *testing.T) { <-done if invoked.Load() != 0 { - t.Fatalf("onCancel fired %d times for an unsignaled task; want 0", + t.Fatalf("onCancel fired %d times for an unsignaled session; want 0", invoked.Load()) } } func TestRequestCancel_EmptyValueStillFires(t *testing.T) { - // Python's task_service.py writes "x" as the value, but a buggy - // caller that wrote "" should not silently keep the watcher + // A caller that wrote "" should not silently keep the watcher // waiting. WatchCancel's contract is "non-empty triggers onCancel"; // we rely on RequestCancel to always set "x" so this test is just // a sanity check that the value round-trips. mr := withCancelClient(t) ctx := context.Background() - if err := RequestCancel(ctx, "task_value"); err != nil { + if err := RequestCancel(ctx, "session_value"); err != nil { t.Fatalf("RequestCancel: %v", err) } - got, err := mr.Get("task_value-cancel") + got, err := mr.Get("session_value-cancel") if err != nil { t.Fatalf("mr.Get: %v", err) } if got != "x" { t.Fatalf("cancel key value = %q, want %q", got, "x") } + if ttl := mr.TTL("session_value-cancel"); ttl <= 0 { + t.Fatalf("cancel key TTL = %v, want finite positive TTL", ttl) + } +} + +func TestCancelRequested(t *testing.T) { + withCancelClient(t) + ctx := context.Background() + + requested, err := CancelRequested(ctx, "session-check") + if err != nil || requested { + t.Fatalf("CancelRequested before marker = %v, %v; want false, nil", requested, err) + } + if err := RequestCancel(ctx, "session-check"); err != nil { + t.Fatalf("RequestCancel: %v", err) + } + requested, err = CancelRequested(ctx, "session-check") + if err != nil || !requested { + t.Fatalf("CancelRequested after marker = %v, %v; want true, nil", requested, err) + } } diff --git a/internal/agent/canvas/canvas.go b/internal/agent/canvas/canvas.go index 1442fe99ec..f3becc8586 100644 --- a/internal/agent/canvas/canvas.go +++ b/internal/agent/canvas/canvas.go @@ -38,8 +38,8 @@ var legacyNoOpNames = map[string]bool{ type CanvasState = runtime.CanvasState // NewCanvasState re-exports runtime.NewCanvasState. -func NewCanvasState(runID, taskID string) *CanvasState { - return runtime.NewCanvasState(runID, taskID) +func NewCanvasState(runID, sessionID string) *CanvasState { + return runtime.NewCanvasState(runID, sessionID) } // Canvas is the in-memory DSL representation loaded from a user_canvas row. diff --git a/internal/agent/canvas/compile.go b/internal/agent/canvas/compile.go index 88cdd9e080..6a29622f17 100644 --- a/internal/agent/canvas/compile.go +++ b/internal/agent/canvas/compile.go @@ -59,8 +59,8 @@ type CompileOptions struct { // Workflow.Invoke), this is a compile-time descriptor: Compile cannot // call compose.WithCheckPointID (the option type is wrong for a // GraphCompileOption), so it only records the id on the returned - // CompiledCanvas — the caller threads it to Invoke. Use a stable, - // per-task value (e.g. taskID) so re-running the same task hits the + // CompiledCanvas — the caller threads it to Invoke. Use a stable value (for + // example a session-derived run id) so resuming hits the // same Redis checkpoint (agent:cp:{id}). When empty, // CompiledCanvas.CheckPointID stays empty and the caller must supply // its own id (or omit it for a fresh per-run checkpoint). @@ -111,8 +111,8 @@ func WithInterruptAfter(nodes []string) CompileOption { // WithCheckPointID sets the stable checkpoint id recorded on the returned // CompiledCanvas. Unlike eino's compose.WithCheckPointID (a run-time // Option), this is a compile-time descriptor: Compile stores the id so the -// caller can pass it to Workflow.Invoke. Pass a stable, per-task value -// (e.g. taskID) so re-running the same task loads the same Redis +// caller can pass it to Workflow.Invoke. Pass a stable, session-derived id +// so resuming loads the same Redis // checkpoint (agent:cp:{id}). func WithCheckPointID(id string) CompileOption { return func(o *CompileOptions) { o.CheckPointID = id } diff --git a/internal/agent/canvas/parallel_subgraph.go b/internal/agent/canvas/parallel_subgraph.go index 412f72df39..306cfff2ef 100644 --- a/internal/agent/canvas/parallel_subgraph.go +++ b/internal/agent/canvas/parallel_subgraph.go @@ -235,7 +235,7 @@ func buildParallelOuterWorkflow( // and Outputs across items. localState, cloneErr := cloneCanvasState(parentState) if cloneErr != nil || localState == nil { - localState = runtime.NewCanvasState(parentState.RunID, parentState.TaskID) + localState = runtime.NewCanvasState(parentState.RunID, parentState.SessionID) localState.Sys = shallowCopyAnyMap(parentState.Sys) localState.Globals = shallowCopyAnyMap(parentState.Globals) } @@ -309,7 +309,7 @@ func cloneCanvasState(src *CanvasState) (*CanvasState, error) { if err != nil { return nil, err } - dst := NewCanvasState(src.RunID, src.TaskID) + dst := NewCanvasState(src.RunID, src.SessionID) if err := json.Unmarshal(raw, dst); err != nil { return nil, err } diff --git a/internal/agent/canvas/run_tracker.go b/internal/agent/canvas/run_tracker.go index 4ca81d6d1b..bc138e6140 100644 --- a/internal/agent/canvas/run_tracker.go +++ b/internal/agent/canvas/run_tracker.go @@ -20,16 +20,18 @@ // // Status code mapping (stored as int under the "status" field): // -// 0 = running, 1 = succeeded, 2 = failed, 3 = cancelled. +// 0 = running, 1 = succeeded, 2 = failed, 3 = cancelled, 4 = waiting_for_user. package canvas import ( "context" "errors" - redis2 "ragflow/internal/engine/redis" + "strconv" "time" "github.com/redis/go-redis/v9" + + redis2 "ragflow/internal/engine/redis" ) // runKeyPrefix is the Redis Hash key namespace for run metadata. @@ -42,6 +44,7 @@ const ( runStatusSucceeded = "1" runStatusFailed = "2" runStatusCancelled = "3" + runStatusWaiting = "4" ) // runFieldInterruptID is the hash field that holds the eino interrupt id a @@ -50,6 +53,22 @@ const runFieldInterruptID = "interrupt_id" func runKey(runID string) string { return runKeyPrefix + runID } +const activeSessionKeyPrefix = "agent:active-session:" + +const activeSessionLeaseTTL = 30 * time.Second + +func activeSessionKey(sessionID string) string { return activeSessionKeyPrefix + sessionID } + +// ActiveSession is the distributed ordinary-Agent run registration. Token is +// an internal lease owner token; it is never exposed as a business run id. +type ActiveSession struct { + SessionID string + Token string + UserID string + CanvasID string + RunID string +} + // RunTracker manages canvas-run metadata (canvas_id, status, checkpoint // link, resume chain, ...) on a Redis Hash. Operations are explicit — the // eino CheckPointStore does NOT write these fields, so callers (HTTP @@ -152,8 +171,9 @@ func (t *RunTracker) GetInterruptID(ctx context.Context, runID string) (string, return v, true, nil } -// ClearInterruptID removes the persisted interrupt id. Called when a run -// completes (no longer paused) or is cancelled (the checkpoint is wiped). +// ClearInterruptID removes the persisted interrupt id after a resumed run +// completes. Cancelling an unrelated turn does not discard the session's +// UserFillUp checkpoint. func (t *RunTracker) ClearInterruptID(ctx context.Context, runID string) error { if t == nil || t.client == nil { return errors.New("run tracker: redis client not initialized") @@ -189,11 +209,161 @@ func (t *RunTracker) MarkCancelled(ctx context.Context, runID string) error { if t == nil || t.client == nil { return errors.New("run tracker: redis client not initialized") } - return t.client.HSet(ctx, runKey(runID), + key := runKey(runID) + pipe := t.client.Pipeline() + pipe.HSet(ctx, key, "status", runStatusCancelled, "finished_at", time.Now().UnixMilli(), "cancel_requested", 1, - ).Err() + ) + pipe.Expire(ctx, key, t.ttl) + _, err := pipe.Exec(ctx) + return err +} + +// RegisterActiveSession atomically acquires the distributed active-run lease +// for sessionID. A cancel marker found without an active lease belongs to an +// older run and is removed before the new owner is registered. A false result +// means another instance already owns the session. +func (t *RunTracker) RegisterActiveSession(ctx context.Context, active ActiveSession) (bool, error) { + if t == nil || t.client == nil { + return false, errors.New("run tracker: redis client not initialized") + } + const script = ` +if redis.call("EXISTS", KEYS[1]) == 1 then return 0 end +redis.call("DEL", KEYS[2]) +redis.call("HSET", KEYS[1], + "session_id", ARGV[1], "token", ARGV[2], "user_id", ARGV[3], + "canvas_id", ARGV[4], "run_id", ARGV[5], "started_at", ARGV[6]) +redis.call("PEXPIRE", KEYS[1], ARGV[7]) +return 1` + result, err := t.client.Eval(ctx, script, + []string{activeSessionKey(active.SessionID), cancelKey(active.SessionID)}, + active.SessionID, active.Token, active.UserID, active.CanvasID, active.RunID, + strconv.FormatInt(time.Now().UnixMilli(), 10), strconv.FormatInt(t.leaseTTL().Milliseconds(), 10), + ).Int64() + return result == 1, err +} + +// GetActiveSession returns the distributed active run for sessionID. A nil +// result means the session is not running. +func (t *RunTracker) GetActiveSession(ctx context.Context, sessionID string) (*ActiveSession, error) { + if t == nil || t.client == nil { + return nil, errors.New("run tracker: redis client not initialized") + } + fields, err := t.client.HGetAll(ctx, activeSessionKey(sessionID)).Result() + if err != nil || len(fields) == 0 { + return nil, err + } + return &ActiveSession{ + SessionID: fields["session_id"], + Token: fields["token"], + UserID: fields["user_id"], + CanvasID: fields["canvas_id"], + RunID: fields["run_id"], + }, nil +} + +// RefreshActiveSession extends the lease only while token still owns it. +func (t *RunTracker) RefreshActiveSession(ctx context.Context, sessionID, token string) (bool, error) { + if t == nil || t.client == nil { + return false, errors.New("run tracker: redis client not initialized") + } + const script = `if redis.call("HGET", KEYS[1], "token") == ARGV[1] then return redis.call("PEXPIRE", KEYS[1], ARGV[2]) end return 0` + result, err := t.client.Eval(ctx, script, []string{activeSessionKey(sessionID)}, token, + strconv.FormatInt(t.leaseTTL().Milliseconds(), 10)).Int64() + return result == 1, err +} + +// ReleaseActiveSession deletes the registration and cancel marker only when +// token still owns the session, preventing stale cleanup from touching a newer +// run. +func (t *RunTracker) ReleaseActiveSession(ctx context.Context, sessionID, token string) (bool, error) { + if t == nil || t.client == nil { + return false, errors.New("run tracker: redis client not initialized") + } + const script = ` +if redis.call("HGET", KEYS[1], "token") ~= ARGV[1] then return 0 end +redis.call("DEL", KEYS[2]) +redis.call("DEL", KEYS[1]) +return 1` + result, err := t.client.Eval(ctx, script, + []string{activeSessionKey(sessionID), cancelKey(sessionID)}, token).Int64() + return result == 1, err +} + +// RequestCancelActiveSession publishes a cancel marker only while token still +// owns the active session lease. This prevents a stale owner from cancelling +// a newer run after its lease has been replaced or released. +func (t *RunTracker) RequestCancelActiveSession(ctx context.Context, sessionID, token string) (bool, error) { + if t == nil || t.client == nil { + return false, errors.New("run tracker: redis client not initialized") + } + const script = ` +if redis.call("HGET", KEYS[1], "token") ~= ARGV[1] then return 0 end +redis.call("SET", KEYS[2], "x", "PX", ARGV[2]) +return 1` + result, err := t.client.Eval(ctx, script, + []string{activeSessionKey(sessionID), cancelKey(sessionID)}, token, + strconv.FormatInt(RequestCancelTTL.Milliseconds(), 10)).Int64() + return result == 1, err +} + +// WatchActiveSession keeps the distributed lease alive for the lifetime of +// ctx. It invokes onLost if another owner replaces the lease or Redis cannot +// confirm ownership before the current lease can expire. +func (t *RunTracker) WatchActiveSession(ctx context.Context, sessionID, token string, onLost func()) { + if t == nil || t.client == nil { + return + } + leaseTTL := t.leaseTTL() + interval := leaseTTL / 3 + if interval > 10*time.Second { + interval = 10 * time.Second + } + if interval < 10*time.Millisecond { + interval = 10 * time.Millisecond + } + ticker := time.NewTicker(interval) + defer ticker.Stop() + lastConfirmed := time.Now() + for { + select { + case <-ctx.Done(): + return + case <-ticker.C: + owned, err := t.RefreshActiveSession(ctx, sessionID, token) + switch { + case err == nil && owned: + lastConfirmed = time.Now() + case err == nil && !owned, time.Since(lastConfirmed) >= leaseTTL*2/3: + if onLost != nil { + onLost() + } + return + } + } + } +} + +func (t *RunTracker) leaseTTL() time.Duration { + if t.ttl > 0 && t.ttl < activeSessionLeaseTTL { + return t.ttl + } + return activeSessionLeaseTTL +} + +// MarkWaiting records a normal UserFillUp pause. It is not a failure. +func (t *RunTracker) MarkWaiting(ctx context.Context, runID string) error { + if t == nil || t.client == nil { + return errors.New("run tracker: redis client not initialized") + } + key := runKey(runID) + pipe := t.client.Pipeline() + pipe.HSet(ctx, key, "status", runStatusWaiting) + pipe.Expire(ctx, key, t.ttl) + _, err := pipe.Exec(ctx) + return err } // Get returns all hash fields for a run. The empty map (not nil) plus a diff --git a/internal/agent/canvas/run_tracker_test.go b/internal/agent/canvas/run_tracker_test.go index 538220f4ec..35ffd7d142 100644 --- a/internal/agent/canvas/run_tracker_test.go +++ b/internal/agent/canvas/run_tracker_test.go @@ -164,6 +164,23 @@ func TestRunTracker_TTLRefresh(t *testing.T) { if len(got) == 0 { t.Fatal("run key expired before refreshed TTL elapsed") } + + // A wait-for-user pause must also keep the run metadata alive for the + // full tracker TTL so the next request can resume it. + mr.FastForward(500 * time.Millisecond) + if err := tracker.MarkWaiting(ctx, "run_ttl"); err != nil { + t.Fatalf("MarkWaiting: %v", err) + } + if d := mr.TTL(runKey("run_ttl")); d < 1500*time.Millisecond { + t.Fatalf("TTL after MarkWaiting = %v, want >= 1.5s", d) + } + got, err = tracker.Get(ctx, "run_ttl") + if err != nil { + t.Fatalf("Get after MarkWaiting: %v", err) + } + if got["status"] != runStatusWaiting { + t.Fatalf("status after MarkWaiting = %q, want %q", got["status"], runStatusWaiting) + } } func TestRunTracker_NilClient(t *testing.T) { @@ -184,7 +201,145 @@ func TestRunTracker_NilClient(t *testing.T) { if err := tracker.MarkCancelled(ctx, "x"); err == nil { t.Fatal("MarkCancelled with nil client: err = nil, want error") } + if err := tracker.MarkWaiting(ctx, "x"); err == nil { + t.Fatal("MarkWaiting with nil client: err = nil, want error") + } if _, err := tracker.Get(ctx, "x"); err == nil { t.Fatal("Get with nil client: err = nil, want error") } } + +func TestRunTrackerActiveSessionRegistration(t *testing.T) { + tracker, mr := newTestTracker(t, time.Hour) + ctx := context.Background() + if err := tracker.client.Set(ctx, cancelKey("session-1"), "x", time.Hour).Err(); err != nil { + t.Fatalf("seed cancel marker: %v", err) + } + active := ActiveSession{ + SessionID: "session-1", Token: "owner-a", UserID: "user-1", + CanvasID: "canvas-1", RunID: "run-1", + } + registered, err := tracker.RegisterActiveSession(ctx, active) + if err != nil || !registered { + t.Fatalf("RegisterActiveSession = %v, %v; want true, nil", registered, err) + } + if mr.Exists(cancelKey("session-1")) { + t.Fatal("registration preserved a cancellation marker from an older run") + } + got, err := tracker.GetActiveSession(ctx, "session-1") + if err != nil || got == nil { + t.Fatalf("GetActiveSession = %#v, %v", got, err) + } + if got.UserID != "user-1" || got.CanvasID != "canvas-1" || got.RunID != "run-1" || got.Token != "owner-a" { + t.Fatalf("unexpected active session: %#v", got) + } + if ttl := mr.TTL(activeSessionKey("session-1")); ttl <= 0 { + t.Fatalf("active session TTL = %v, want finite positive TTL", ttl) + } + registered, err = tracker.RegisterActiveSession(ctx, ActiveSession{SessionID: "session-1", Token: "owner-b"}) + if err != nil || registered { + t.Fatalf("second RegisterActiveSession = %v, %v; want false, nil", registered, err) + } + requested, err := tracker.RequestCancelActiveSession(ctx, "session-1", "owner-a") + if err != nil || !requested { + t.Fatalf("RequestCancelActiveSession = %v, %v; want true, nil", requested, err) + } + registered, err = tracker.RegisterActiveSession(ctx, ActiveSession{SessionID: "session-1", Token: "owner-b"}) + if err != nil || registered { + t.Fatalf("registration racing active cancel = %v, %v; want false, nil", registered, err) + } + if !mr.Exists(cancelKey("session-1")) { + t.Fatal("failed registration erased the active owner's cancel marker") + } + registered, err = tracker.RegisterActiveSession(ctx, ActiveSession{SessionID: "session-2", Token: "owner-b"}) + if err != nil || !registered { + t.Fatalf("different-session RegisterActiveSession = %v, %v; want true, nil", registered, err) + } +} + +func TestRunTrackerActiveSessionTokenProtectsRefreshAndRelease(t *testing.T) { + tracker, _ := newTestTracker(t, time.Hour) + ctx := context.Background() + registered, err := tracker.RegisterActiveSession(ctx, ActiveSession{SessionID: "session-1", Token: "owner-a"}) + if err != nil || !registered { + t.Fatalf("RegisterActiveSession = %v, %v", registered, err) + } + owned, err := tracker.RefreshActiveSession(ctx, "session-1", "owner-b") + if err != nil || owned { + t.Fatalf("foreign RefreshActiveSession = %v, %v; want false, nil", owned, err) + } + released, err := tracker.ReleaseActiveSession(ctx, "session-1", "owner-b") + if err != nil || released { + t.Fatalf("foreign ReleaseActiveSession = %v, %v; want false, nil", released, err) + } + if got, _ := tracker.GetActiveSession(ctx, "session-1"); got == nil { + t.Fatal("foreign release removed active session") + } + if err := tracker.client.Set(ctx, cancelKey("session-1"), "x", time.Hour).Err(); err != nil { + t.Fatalf("seed cancel marker: %v", err) + } + released, err = tracker.ReleaseActiveSession(ctx, "session-1", "owner-a") + if err != nil || !released { + t.Fatalf("owner ReleaseActiveSession = %v, %v; want true, nil", released, err) + } + if got, _ := tracker.GetActiveSession(ctx, "session-1"); got != nil { + t.Fatalf("active session remains after release: %#v", got) + } + if tracker.client.Exists(ctx, cancelKey("session-1")).Val() != 0 { + t.Fatal("release did not clear session cancel marker") + } +} + +func TestRunTrackerActiveSessionTokenProtectsCancel(t *testing.T) { + tracker, _ := newTestTracker(t, time.Hour) + ctx := context.Background() + registered, err := tracker.RegisterActiveSession(ctx, ActiveSession{SessionID: "session-1", Token: "owner-a"}) + if err != nil || !registered { + t.Fatalf("RegisterActiveSession = %v, %v", registered, err) + } + requested, err := tracker.RequestCancelActiveSession(ctx, "session-1", "owner-b") + if err != nil || requested { + t.Fatalf("foreign RequestCancelActiveSession = %v, %v; want false, nil", requested, err) + } + if tracker.client.Exists(ctx, cancelKey("session-1")).Val() != 0 { + t.Fatal("foreign cancel created a marker") + } + requested, err = tracker.RequestCancelActiveSession(ctx, "session-1", "owner-a") + if err != nil || !requested { + t.Fatalf("owner RequestCancelActiveSession = %v, %v; want true, nil", requested, err) + } + if tracker.client.Exists(ctx, cancelKey("session-1")).Val() != 1 { + t.Fatal("owner cancel did not create a marker") + } + released, err := tracker.ReleaseActiveSession(ctx, "session-1", "owner-a") + if err != nil || !released { + t.Fatalf("owner ReleaseActiveSession = %v, %v; want true, nil", released, err) + } + requested, err = tracker.RequestCancelActiveSession(ctx, "session-1", "owner-a") + if err != nil || requested { + t.Fatalf("late RequestCancelActiveSession = %v, %v; want false, nil", requested, err) + } + if tracker.client.Exists(ctx, cancelKey("session-1")).Val() != 0 { + t.Fatal("late cancel recreated a marker after active-session release") + } +} + +func TestRunTrackerWatchActiveSessionCancelsWhenLeaseOwnershipIsLost(t *testing.T) { + tracker, _ := newTestTracker(t, 90*time.Millisecond) + ctx, cancel := context.WithCancel(t.Context()) + defer cancel() + registered, err := tracker.RegisterActiveSession(ctx, ActiveSession{SessionID: "session-1", Token: "owner-a"}) + if err != nil || !registered { + t.Fatalf("RegisterActiveSession = %v, %v", registered, err) + } + lost := make(chan struct{}) + go tracker.WatchActiveSession(ctx, "session-1", "owner-a", func() { close(lost) }) + if err := tracker.client.HSet(ctx, activeSessionKey("session-1"), "token", "owner-b").Err(); err != nil { + t.Fatalf("replace lease owner: %v", err) + } + select { + case <-lost: + case <-time.After(time.Second): + t.Fatal("lease watcher did not cancel after ownership changed") + } +} diff --git a/internal/agent/canvas/runner.go b/internal/agent/canvas/runner.go index 003fa8fdb3..cc022d866b 100644 --- a/internal/agent/canvas/runner.go +++ b/internal/agent/canvas/runner.go @@ -66,19 +66,20 @@ import ( // The handler converts each RunEvent into one SSE frame in the // Python-shaped envelope: // -// data:{"event":"","message_id":"","created_at":,"task_id":"","session_id":"","data":} +// data:{"event":"","message_id":"","created_at":,"session_id":"","data":} // // Type is the event tag; Data is the JSON payload string (already // serialised — handler does not re-marshal). The handler wraps Data // into the "data" field of the outer envelope so the front-end's // use-send-message.ts parser sees a flat {event, message_id, -// created_at, task_id, session_id, data} object on every frame. +// created_at, session_id, data} object on every frame. +// WriteChatbotRunEvent may additionally expose task_id=session_id as a wire +// alias for existing clients; RunEvent itself has only one run identity. type RunEvent struct { Type string Data string MessageID string CreatedAt int64 - TaskID string SessionID string } @@ -151,18 +152,15 @@ type ErrorEvent struct { // - any other non-nil error: run failed; surface as `error` event. type RunFunc func(ctx context.Context, root map[string]any) (*CanvasState, error) -// Runner is the per-canvas execution runtime. It owns the +// Runner is the ordinary-Agent execution runtime. It owns the // interrupt-id map (V1 in-memory persistence keyed by -// (canvasID, sessionID)) and the goroutine cancellation registry. +// (canvasID, sessionID)). The service owns the run context and cancellation. // // Concurrency: Runner methods are safe for concurrent use. The -// output channel is owned by the goroutine that started a run; -// the Cancel method signals the underlying run via the cancel -// channel that the RunFunc is expected to observe. +// output channel is owned by the goroutine that started a run. type Runner struct { mu sync.Mutex interruptIDs map[string]string // key = canvasID + "|" + sessionID; value = eino interrupt id - runCancels map[string]chan struct{} } // NewRunner returns a fresh Runner with the in-memory interrupt-id @@ -171,7 +169,6 @@ type Runner struct { func NewRunner() *Runner { return &Runner{ interruptIDs: make(map[string]string), - runCancels: make(map[string]chan struct{}), } } @@ -213,8 +210,8 @@ func (r *Runner) getInterruptID(canvasID, sessionID string) string { // four-outcome flow. The channel is always closed on return so the // handler's for-range loop terminates. // -// Metadata injection: the output channel, message_id, task_id, and -// session_id are injected into root so the RunFunc (buildRunFunc in +// Metadata injection: the output channel, message_id, and session_id are +// injected into root so the RunFunc (buildRunFunc in // service/agent.go) can emit intermediate events (workflow_started, // node_started, node_finished, workflow_finished) during execution // rather than only after the invoke completes. The key names follow @@ -230,53 +227,23 @@ func (r *Runner) Run( out := make(chan RunEvent, 8) if run == nil { - pushErr(out, "canvas: nil RunFunc") + pushErr(out, "canvas: nil RunFunc", sessionID) close(out) return out } - cancel := make(chan struct{}) - r.mu.Lock() - if prev, hadPrev := r.runCancels[canvasID]; hadPrev { - select { - case <-prev: - default: - close(prev) - } - } - r.runCancels[canvasID] = cancel - r.mu.Unlock() - - // Generate the identifiers the RunFunc and SSE envelope need. - // message_id is generated per-run so the front-end can correlate - // all events for a single user turn. task_id is the published - // version id (if available) or a per-run UUID. + // Generate the message identifier the RunFunc and SSE envelope need. messageID := utility.GenerateToken() - taskID := "" - if v, ok := root["version_id"].(string); ok && v != "" { - taskID = v - } - if taskID == "" { - taskID = utility.GenerateToken() - } // Inject the output channel + metadata so the RunFunc can emit // events during execution (workflow_started, node_started, // node_finished, etc.). root["__events__"] = out root["__message_id__"] = messageID - root["__task_id__"] = taskID root["__session_id__"] = sessionID go func() { defer close(out) - defer func() { - r.mu.Lock() - if r.runCancels[canvasID] == cancel { - delete(r.runCancels, canvasID) - } - r.mu.Unlock() - }() // Panic sentinel (temporary diagnostic — see plan): // a panic anywhere in the run goroutine used to silently // propagate, leaving the events channel closed-empty so the @@ -299,15 +266,16 @@ func (r *Runner) Run( // root inside the RunFunc — see service/agent.go's // buildRunFunc. if userInput != nil { - if id := r.getInterruptID(canvasID, sessionID); id != "" { + id := r.getInterruptID(canvasID, sessionID) + if _, hasPersistedID := root["__resume_interrupt_id__"]; !hasPersistedID && id != "" { root["__resume_interrupt_id__"] = id root["__resume_data__"] = userInput } } - _, runErr := safeInvoke(ctx, cancel, run, root) + _, runErr := safeInvoke(ctx, run, root) if runErr != nil { - if errors.Is(runErr, context.Canceled) || errors.Is(runErr, errCancelled) { + if errors.Is(runErr, context.Canceled) { return } if ctxs := ExtractInterruptContexts(runErr); len(ctxs) > 0 { @@ -320,7 +288,6 @@ func (r *Runner) Run( common.Info("canvas runner interrupt", zap.String("canvas", canvasID), zap.String("session", sessionID), - zap.String("task", taskID), zap.String("contexts", formatInterruptContexts(ctxs)), zap.String("display", displayID), zap.String("resume", resumeID)) @@ -336,7 +303,7 @@ func (r *Runner) Run( } } } - push(out, RunEvent{Type: "waiting_for_user", Data: safeEventJSON(waiting), MessageID: messageID, CreatedAt: nowUnix(), TaskID: taskID, SessionID: sessionID}) + push(out, RunEvent{Type: "waiting_for_user", Data: safeEventJSON(waiting), MessageID: messageID, CreatedAt: nowUnix(), SessionID: sessionID}) return } if IsInterruptError(runErr) { @@ -345,10 +312,10 @@ func (r *Runner) Run( // without a cpn id — the front-end falls back to // the first paused session it knows about. r.saveInterruptID(canvasID, sessionID, runErr.Error()) - push(out, RunEvent{Type: "waiting_for_user", Data: safeEventJSON(WaitingForUserEvent{CpnID: runErr.Error()}), MessageID: messageID, CreatedAt: nowUnix(), TaskID: taskID, SessionID: sessionID}) + push(out, RunEvent{Type: "waiting_for_user", Data: safeEventJSON(WaitingForUserEvent{CpnID: runErr.Error()}), MessageID: messageID, CreatedAt: nowUnix(), SessionID: sessionID}) return } - pushErr(out, runErr.Error()) + pushErr(out, runErr.Error(), sessionID) return } }() @@ -356,22 +323,6 @@ func (r *Runner) Run( return out } -// Cancel signals an in-flight run for the given canvas to stop. -// Safe to call when no run is active. -func (r *Runner) Cancel(canvasID string) { - r.mu.Lock() - cancel, ok := r.runCancels[canvasID] - r.mu.Unlock() - if !ok { - return - } - select { - case <-cancel: - default: - close(cancel) - } -} - // Peek reports whether a paused interrupt id is held for the given // (canvasID, sessionID). It is intended for tests and diagnostics; // the real runner does not need it at run time. @@ -382,16 +333,9 @@ func (r *Runner) Peek(canvasID, sessionID string) bool { return ok } -// errCancelled is the sentinel safeInvoke returns when the cancel -// channel fires during a run. It is wrapped against context.Canceled -// so callers can `errors.Is` either. -var errCancelled = fmt.Errorf("canvas: run cancelled") - -// safeInvoke calls the supplied RunFunc with context-cancel and -// driver-cancel both wired in. The RunFunc is expected to honour -// ctx.Done() — the cancel channel is a secondary signal for the -// V1 in-process driver. -func safeInvoke(ctx context.Context, cancel chan struct{}, run RunFunc, root map[string]any) (*CanvasState, error) { +// safeInvoke calls the supplied RunFunc with the managed child context. +// The RunFunc is expected to honour ctx.Done(). +func safeInvoke(ctx context.Context, run RunFunc, root map[string]any) (*CanvasState, error) { done := make(chan struct{}) var ( state *CanvasState @@ -415,9 +359,16 @@ func safeInvoke(ctx context.Context, cancel chan struct{}, run RunFunc, root map }() select { case <-done: + if ctx.Err() != nil { + return nil, ctx.Err() + } return state, err - case <-cancel: - return nil, errCancelled + case <-ctx.Done(): + // Do not abandon the workflow goroutine. Eino and context-aware + // HTTP/tool calls should return promptly once the child context is + // cancelled; waiting here keeps the run under management. + <-done + return nil, ctx.Err() } } @@ -440,7 +391,7 @@ func push(out chan<- RunEvent, ev RunEvent) { } // pushErr serialises an ErrorEvent and pushes it on the channel. -func pushErr(out chan<- RunEvent, msg string) { +func pushErr(out chan<- RunEvent, msg, sessionID string) { payload, err := json.Marshal(ErrorEvent{Message: msg}) if err != nil { common.Warn("runner: pushErr json.Marshal failed, falling back", @@ -449,7 +400,7 @@ func pushErr(out chan<- RunEvent, msg string) { // Fall back to a hard-coded minimal JSON. payload = []byte(`{"message":"event serialization failed"}`) } - push(out, RunEvent{Type: "error", Data: string(payload)}) + push(out, RunEvent{Type: "error", Data: string(payload), SessionID: sessionID, CreatedAt: nowUnix()}) } // safeEventJSON marshals v to a JSON string, falling back to diff --git a/internal/agent/canvas/runner_cancel_test.go b/internal/agent/canvas/runner_cancel_test.go new file mode 100644 index 0000000000..6cc00af7d0 --- /dev/null +++ b/internal/agent/canvas/runner_cancel_test.go @@ -0,0 +1,68 @@ +// Copyright 2026 The InfiniFlow Authors. All Rights Reserved. +// +// Licensed under the Apache License, Version 2.0 (the "License"); + +package canvas + +import ( + "context" + "testing" + "time" +) + +func blockingRun(started chan<- struct{}) RunFunc { + return func(ctx context.Context, _ map[string]any) (*CanvasState, error) { + close(started) + <-ctx.Done() + return nil, ctx.Err() + } +} + +func waitClosed(t *testing.T, events <-chan RunEvent) { + t.Helper() + select { + case _, ok := <-events: + if ok { + for range events { + } + } + case <-time.After(2 * time.Second): + t.Fatal("run event channel did not close after cancellation") + } +} + +func TestRunnerUsesSessionMetadata(t *testing.T) { + r := NewRunner() + root := map[string]any{} + started := make(chan struct{}) + ctx, cancel := context.WithCancel(t.Context()) + events := r.Run(ctx, blockingRun(started), "canvas-1", "session-1", nil, root) + <-started + if got := root["__session_id__"]; got != "session-1" { + t.Fatalf("session metadata = %v, want session-1", got) + } + cancel() + waitClosed(t, events) +} + +func TestRunnerParentContextCancelsManagedRun(t *testing.T) { + r := NewRunner() + ctx, cancel := context.WithCancel(context.Background()) + started := make(chan struct{}) + returned := make(chan struct{}) + run := func(ctx context.Context, _ map[string]any) (*CanvasState, error) { + close(started) + <-ctx.Done() + close(returned) + return nil, ctx.Err() + } + events := r.Run(ctx, run, "canvas", "session", nil, map[string]any{}) + <-started + cancel() + waitClosed(t, events) + select { + case <-returned: + case <-time.After(time.Second): + t.Fatal("RunFunc was left running after parent context cancellation") + } +} diff --git a/internal/agent/canvas/scheduler.go b/internal/agent/canvas/scheduler.go index 19f3b1d20a..452d1dc9b5 100644 --- a/internal/agent/canvas/scheduler.go +++ b/internal/agent/canvas/scheduler.go @@ -34,7 +34,7 @@ import ( ) // ctxKey is the unexported context-key type for per-run metadata -// (events channel, message/task/session ids) so the statePre/statePost +// (events channel, message/session ids) so the statePre/statePost // wrappers can emit node_started/node_finished without depending on // the service package. type ctxKey string @@ -46,7 +46,6 @@ const terminalMergeNodeID = "__canvas_terminal_merge__" type RunMeta struct { Events chan RunEvent MessageID string - TaskID string SessionID string } @@ -274,7 +273,7 @@ func sanitizeNodeInputs(inputs map[string]any) map[string]any { // nodeStartedAt records the per-node start time in state.Sys and emits a // node_started RunEvent. Called from the per-node statePre wrapper. -// Metadata (message/task/session ids) is read from ctx via RunMeta. +// Metadata (message/session ids) is read from ctx via RunMeta. func nodeStartedAt(ctx context.Context, state *CanvasState, cpnID, componentName, componentType string, inputs map[string]any) { common.Debug("node_started", zap.String("cpnID", cpnID), zap.String("componentName", componentName)) if state == nil { @@ -295,14 +294,14 @@ func nodeStartedAt(ctx context.Context, state *CanvasState, cpnID, componentName Thoughts: "", }) meta := GetRunMeta(ctx) - msgID, taskID, sessionID := "", "", "" + msgID, sessionID := "", "" if meta != nil { - msgID, taskID, sessionID = meta.MessageID, meta.TaskID, meta.SessionID + msgID, sessionID = meta.MessageID, meta.SessionID } emitEventFromCtx(ctx, RunEvent{ Type: "node_started", Data: string(nsData), MessageID: msgID, CreatedAt: time.Now().Unix(), - TaskID: taskID, SessionID: sessionID, + SessionID: sessionID, }) } @@ -358,14 +357,14 @@ func nodeFinishedNow(ctx context.Context, state *CanvasState, cpnID, componentNa CreatedAt: now, }) meta := GetRunMeta(ctx) - msgID, taskID, sessionID := "", "", "" + msgID, sessionID := "", "" if meta != nil { - msgID, taskID, sessionID = meta.MessageID, meta.TaskID, meta.SessionID + msgID, sessionID = meta.MessageID, meta.SessionID } emitEventFromCtx(ctx, RunEvent{ Type: "node_finished", Data: string(nfData), MessageID: msgID, CreatedAt: time.Now().Unix(), - TaskID: taskID, SessionID: sessionID, + SessionID: sessionID, }) } @@ -412,7 +411,7 @@ func BuildWorkflow(ctx context.Context, c *Canvas) (*compose.Workflow[map[string globals := c.Globals genState := func(runCtx context.Context) *CanvasState { if ctxState, _, _ := runtime.GetStateFromContext[*runtime.CanvasState](runCtx); ctxState != nil { - st := NewCanvasState(ctxState.RunID, ctxState.TaskID) + st := NewCanvasState(ctxState.RunID, ctxState.SessionID) for cpnID, bucket := range ctxState.Snapshot() { for key, value := range bucket { st.SetVar(cpnID, key, value) diff --git a/internal/agent/canvas/state_serializer_test.go b/internal/agent/canvas/state_serializer_test.go index f9c6f47e55..850c989b91 100644 --- a/internal/agent/canvas/state_serializer_test.go +++ b/internal/agent/canvas/state_serializer_test.go @@ -23,7 +23,7 @@ import ( ) func TestCanvasStateSerializer_RoundTrip(t *testing.T) { - src := NewCanvasState("run_abc", "task_xyz") + src := NewCanvasState("run_abc", "session_xyz") src.Outputs["retrieval_0"] = map[string]any{ "chunks": []string{"a", "b", "c"}, "doc_aggs": map[string]int{"doc1": 3, "doc2": 1}, @@ -64,8 +64,8 @@ func TestCanvasStateSerializer_RoundTrip(t *testing.T) { if dst.RunID != src.RunID { t.Fatalf("RunID = %q, want %q", dst.RunID, src.RunID) } - if dst.TaskID != src.TaskID { - t.Fatalf("TaskID = %q, want %q", dst.TaskID, src.TaskID) + if dst.SessionID != src.SessionID { + t.Fatalf("SessionID = %q, want %q", dst.SessionID, src.SessionID) } // JSON round-trip coerces numbers to float64, so we re-marshal both // sides and compare bytes — that is the real contract of the @@ -110,8 +110,8 @@ func TestCanvasStateSerializer_EmptyState(t *testing.T) { if err := ser.Unmarshal(data, dst); err != nil { t.Fatalf("Unmarshal empty: %v", err) } - if dst.RunID != "r" || dst.TaskID != "t" { - t.Fatalf("ids not preserved: %q %q", dst.RunID, dst.TaskID) + if dst.RunID != "r" || dst.SessionID != "t" { + t.Fatalf("ids not preserved: %q %q", dst.RunID, dst.SessionID) } } @@ -135,8 +135,8 @@ func TestCanvasStateSerializer_UnmarshalIntoExistingPointer(t *testing.T) { t.Fatalf("Sys[x] = %v (%T), want float64(1)", dst.Sys["x"], dst.Sys["x"]) } // Ids are overwritten by the round-trip. - if dst.RunID != "r2" || dst.TaskID != "t2" { - t.Fatalf("ids not overwritten: %q %q", dst.RunID, dst.TaskID) + if dst.RunID != "r2" || dst.SessionID != "t2" { + t.Fatalf("ids not overwritten: %q %q", dst.RunID, dst.SessionID) } } diff --git a/internal/agent/canvas/stream.go b/internal/agent/canvas/stream.go index d6d3be72d0..8ad7c86447 100644 --- a/internal/agent/canvas/stream.go +++ b/internal/agent/canvas/stream.go @@ -32,8 +32,8 @@ import ( type StreamEvent struct { // Event is the event name: "node_start" | "node_finish" | "message" | "error" | "cancelled" | ... Event string `json:"event"` - // TaskID identifies the canvas run; required for client correlation. - TaskID string `json:"task_id"` + // SessionID identifies the canvas session; required for client correlation. + SessionID string `json:"session_id"` // Component identifies the canvas component that produced the event. Component string `json:"component,omitempty"` // Data is the free-form event body. SSE wire format is "data: " + json(ev.Data). @@ -74,7 +74,7 @@ func (e *channelEmitter) Emit(ev StreamEvent) error { default: common.Warn("canvas stream: dropping event (buffer full)", zap.String("event", ev.Event), - zap.String("task", ev.TaskID)) + zap.String("session", ev.SessionID)) return nil } } diff --git a/internal/agent/canvas/stream_test.go b/internal/agent/canvas/stream_test.go index 97e8ecd3f3..be9ba55449 100644 --- a/internal/agent/canvas/stream_test.go +++ b/internal/agent/canvas/stream_test.go @@ -28,10 +28,10 @@ func TestChannelEmitter_EmitAndClose(t *testing.T) { ch := em.(*channelEmitter).Channel() evs := []StreamEvent{ - {Event: "node_start", TaskID: "t1", Component: "begin_0"}, - {Event: "message", TaskID: "t1", Component: "llm_0", + {Event: "node_start", SessionID: "s1", Component: "begin_0"}, + {Event: "message", SessionID: "s1", Component: "llm_0", Data: map[string]any{"delta": "hello"}}, - {Event: "node_finish", TaskID: "t1", Component: "begin_0", + {Event: "node_finish", SessionID: "s1", Component: "begin_0", Data: map[string]any{"ok": true}}, } for _, ev := range evs { @@ -51,7 +51,7 @@ func TestChannelEmitter_EmitAndClose(t *testing.T) { t.Fatalf("got %d events, want %d", len(got), len(evs)) } for i, ev := range got { - if ev.Event != evs[i].Event || ev.TaskID != evs[i].TaskID || + if ev.Event != evs[i].Event || ev.SessionID != evs[i].SessionID || ev.Component != evs[i].Component { t.Fatalf("event %d: got %+v, want %+v", i, ev, evs[i]) } @@ -62,12 +62,12 @@ func TestChannelEmitter_NonBlockingDrop(t *testing.T) { // Buffer of 1 with no reader; the second Emit must return nil // immediately (drop on full) rather than block. em := NewChannelEmitter(1) - if err := em.Emit(StreamEvent{Event: "e1", TaskID: "t"}); err != nil { + if err := em.Emit(StreamEvent{Event: "e1", SessionID: "s"}); err != nil { t.Fatalf("Emit 1: %v", err) } done := make(chan struct{}) go func() { - if err := em.Emit(StreamEvent{Event: "e2", TaskID: "t"}); err != nil { + if err := em.Emit(StreamEvent{Event: "e2", SessionID: "s"}); err != nil { t.Errorf("Emit 2: %v", err) } close(done) @@ -88,7 +88,7 @@ func TestChannelEmitter_NonBlockingDrop(t *testing.T) { func TestFormatSSE(t *testing.T) { ev := StreamEvent{ Event: "message", - TaskID: "task_42", + SessionID: "session_42", Component: "llm_0", Data: map[string]any{ "delta": "héllo, 世界", @@ -121,7 +121,7 @@ func TestFormatSSE(t *testing.T) { func TestFormatSSE_EmptyData(t *testing.T) { // Empty Data must still produce a valid frame, not panic. - got := FormatSSE(StreamEvent{Event: "node_start", TaskID: "t"}) + got := FormatSSE(StreamEvent{Event: "node_start", SessionID: "s"}) if !strings.HasPrefix(got, "data: ") || !strings.HasSuffix(got, "\n\n") { t.Fatalf("empty Data frame malformed: %q", got) } diff --git a/internal/agent/component/agent.go b/internal/agent/component/agent.go index 4831570796..1ba3fe3e0a 100644 --- a/internal/agent/component/agent.go +++ b/internal/agent/component/agent.go @@ -161,7 +161,7 @@ func runEinoReActAgent(ctx context.Context, p AgentParam) (*schema.Message, erro if err != nil { return nil, fmt.Errorf("build model: %w", err) } - tools, err := buildAgentTools(p) + tools, err := buildAgentTools(ctx, p) if err != nil { return nil, fmt.Errorf("build tools: %w", err) } @@ -420,13 +420,16 @@ func chunksFromState(ctx context.Context) []prompts.CitationSource { // are skipped silently. func (c *AgentComponent) GetInputForm() map[string]any { out := extractAgentPromptInputForm(c.param.SystemPrompt, c.param.UserPrompt) - tools, err := buildAgentTools(c.param) + // GetInputForm is metadata introspection outside a Canvas run and its + // interface has no context. Runtime tool construction uses the run ctx in + // runEinoReActAgent above. + metadataCtx := context.Background() + tools, err := buildAgentTools(metadataCtx, c.param) if err != nil { return out } - ctx := context.Background() for _, t := range tools { - info, ierr := t.Info(ctx) + info, ierr := t.Info(metadataCtx) name := "" if ierr == nil && info != nil { name = info.Name @@ -480,7 +483,9 @@ func sortedAgentPromptInputKeys(systemPrompt, userPrompt string) []string { // interface. Mirrors Python's per-tool reset() — useful for clearing // per-invocation state (caches, scratch buffers) between calls. func (c *AgentComponent) Reset() { - tools, err := buildAgentTools(c.param) + // Reset is a context-free lifecycle hook and does not execute a tool. + // Runtime tool construction receives the Canvas run context. + tools, err := buildAgentTools(context.Background(), c.param) if err != nil { return } @@ -546,14 +551,14 @@ func optimizeMultiTurnQuestion(ctx context.Context, db *gorm.DB, p AgentParam, h return strings.TrimSpace(resp.Content), nil } -func buildAgentTools(p AgentParam) ([]einotool.BaseTool, error) { +func buildAgentTools(ctx context.Context, p AgentParam) ([]einotool.BaseTool, error) { tools, err := agenttool.BuildAll(p.Tools, p.ToolParams) if err != nil { return nil, err } toolNames := make(map[string]struct{}, len(tools)+len(p.SubAgents)) for _, tool := range tools { - info, err := tool.Info(context.Background()) + info, err := tool.Info(ctx) if err != nil { return nil, fmt.Errorf("agent tool info: %w", err) } diff --git a/internal/agent/component/agent_test.go b/internal/agent/component/agent_test.go index dc90c2e6db..9f5246e3e2 100644 --- a/internal/agent/component/agent_test.go +++ b/internal/agent/component/agent_test.go @@ -649,7 +649,7 @@ func TestAgent_NewAcceptsCanvasToolObjects(t *testing.T) { if agent.param.ToolParams["retrieval"] == nil { t.Fatalf("agent.param.ToolParams missing retrieval: %#v", agent.param.ToolParams) } - if _, err := buildAgentTools(agent.param); err != nil { + if _, err := buildAgentTools(t.Context(), agent.param); err != nil { t.Fatalf("buildAgentTools: %v", err) } } @@ -698,7 +698,7 @@ func TestAgent_CanCreateReactAgentWithAllRegisteredTools(t *testing.T) { }, MaxRounds: 1, } - tools, err := buildAgentTools(p) + tools, err := buildAgentTools(t.Context(), p) if err != nil { t.Fatalf("buildAgentTools: %v", err) } @@ -760,7 +760,7 @@ func TestAgent_CanvasSubAgentToolBuildsDynamicTool(t *testing.T) { t.Fatalf("sub agents = %d, want 1", len(agent.param.SubAgents)) } - tools, err := buildAgentTools(agent.param) + tools, err := buildAgentTools(t.Context(), agent.param) if err != nil { t.Fatalf("buildAgentTools: %v", err) } @@ -810,7 +810,7 @@ func TestAgent_CanvasSubAgentToolNamesAreUniqueAfterNormalization(t *testing.T) t.Fatalf("New(Agent) returned %T, want *AgentComponent", c) } - tools, err := buildAgentTools(agent.param) + tools, err := buildAgentTools(t.Context(), agent.param) if err != nil { t.Fatalf("buildAgentTools: %v", err) } diff --git a/internal/agent/component/message.go b/internal/agent/component/message.go index 8b21ea146c..9b4f6d580c 100644 --- a/internal/agent/component/message.go +++ b/internal/agent/component/message.go @@ -316,8 +316,8 @@ func (m *MessageComponent) Invoke(ctx context.Context, db *gorm.DB, inputs map[s saveErr := saver.Save(ctx, MemorySaveRequest{ MemoryIDs: memIDs, UserID: userID, - AgentID: state.TaskID, - SessionID: state.RunID, + AgentID: stringFromStateSys(state, "agent_id"), + SessionID: state.SessionID, UserInput: stringFromStateSys(state, "query"), AgentResponse: rendered, }) diff --git a/internal/agent/component/message_phase8b_test.go b/internal/agent/component/message_phase8b_test.go index cadda9ac29..c437aacd17 100644 --- a/internal/agent/component/message_phase8b_test.go +++ b/internal/agent/component/message_phase8b_test.go @@ -255,8 +255,9 @@ func TestMessage_MemorySave_Success(t *testing.T) { defer SetMemorySaver(nil) c, _ := NewMessageComponent(map[string]any{"text": "hi"}) - state := canvas.NewCanvasState("run-y", "task-y") + state := canvas.NewCanvasState("run-y", "session-y") state.Sys["query"] = "what?" + state.Sys["agent_id"] = "agent-y" ctx := withStateForTest(context.Background(), state) _, err := c.Invoke(ctx, nil, map[string]any{ @@ -271,11 +272,11 @@ func TestMessage_MemorySave_Success(t *testing.T) { if len(saved.MemoryIDs) != 2 || saved.MemoryIDs[0] != "m1" || saved.MemoryIDs[1] != "m2" { t.Errorf("MemoryIDs: %+v", saved.MemoryIDs) } - if saved.AgentID != "task-y" { - t.Errorf("AgentID: got %q, want task-y", saved.AgentID) + if saved.AgentID != "agent-y" { + t.Errorf("AgentID: got %q, want agent-y", saved.AgentID) } - if saved.SessionID != "run-y" { - t.Errorf("SessionID: got %q, want run-y", saved.SessionID) + if saved.SessionID != "session-y" { + t.Errorf("SessionID: got %q, want session-y", saved.SessionID) } if saved.UserInput != "what?" { t.Errorf("UserInput: got %q, want what?", saved.UserInput) diff --git a/internal/agent/component/stagehand_runtime.go b/internal/agent/component/stagehand_runtime.go index 07c01ab360..7214f1f355 100644 --- a/internal/agent/component/stagehand_runtime.go +++ b/internal/agent/component/stagehand_runtime.go @@ -408,7 +408,9 @@ func (r *stagehandRuntime) RunTask(ctx context.Context, req RunTaskRequest) (str // Best-effort End. The deferred call's error is intentionally // dropped — the agent result has already been returned; the // End failure is logged by the stagehand server. - _, _ = client.Sessions.End(context.Background(), sessionID, stagehand.SessionEndParams{}) + cleanupCtx, cancel := context.WithTimeout(context.WithoutCancel(ctx), 10*time.Second) + defer cancel() + _, _ = client.Sessions.End(cleanupCtx, sessionID, stagehand.SessionEndParams{}) }() // Sessions.Execute runs the multi-step agent. MaxSteps defaults @@ -558,7 +560,9 @@ func (r *stagehandRuntime) RunExtract(ctx context.Context, req RunExtractRequest } sessionID := startResp.Data.SessionID defer func() { - _, _ = client.Sessions.End(context.Background(), sessionID, stagehand.SessionEndParams{}) + cleanupCtx, cancel := context.WithTimeout(context.WithoutCancel(ctx), 10*time.Second) + defer cancel() + _, _ = client.Sessions.End(cleanupCtx, sessionID, stagehand.SessionEndParams{}) }() // Optional: navigate to the target URL before extracting. diff --git a/internal/agent/component/tool_dispatch_test.go b/internal/agent/component/tool_dispatch_test.go index 2886611a57..75f16f1253 100644 --- a/internal/agent/component/tool_dispatch_test.go +++ b/internal/agent/component/tool_dispatch_test.go @@ -49,7 +49,7 @@ func TestPhase3_5_ToolDispatchViaEinoReact(t *testing.T) { } // TestPhase3_6_ToolDSLLoading exercises the Python-equivalent of -// _load_tool_obj: buildAgentTools(p) maps AgentParam.Tools (a slice +// _load_tool_obj: buildAgentTools(ctx, p) maps AgentParam.Tools (a slice // of registered tool names) to einotool.BaseTool instances via // agenttool.BuildAll. The factory validates each name against the // registry; an unknown name surfaces a build-time error (not a @@ -102,7 +102,7 @@ func TestAgent_GoogleToolDSLParamsLoading(t *testing.T) { t.Fatalf("google tool form missing q: %+v", googleForm) } - if _, err := buildAgentTools(c.param); err != nil { + if _, err := buildAgentTools(t.Context(), c.param); err != nil { t.Fatalf("buildAgentTools with google params: %v", err) } } diff --git a/internal/agent/dsl/reset.go b/internal/agent/dsl/reset.go index 668cec18f1..382869f777 100644 --- a/internal/agent/dsl/reset.go +++ b/internal/agent/dsl/reset.go @@ -20,7 +20,7 @@ // Python's Canvas.reset() does two things: // // 1. Graph.reset(): clears the per-component state (path, in-memory -// caches) and removes the per-task Redis log/cancel keys. +// caches) and removes the per-session Redis log/cancel keys. // // 2. Per-run state wipe: empties self.history / retrieval / memory, // then walks self.globals to zero out every "sys.*" key and to @@ -30,8 +30,8 @@ // In the Go port there is no per-canvas "Graph" runtime — the // executor is reconstructed from the DSL on every Run. So the // Python "Graph.reset()" side (step 1) is implicitly handled by the -// per-run rebuild and the per-task Redis keys are still owned by the -// Python task executor. The Go port is responsible for the +// per-run rebuild and the Redis keys are still owned by the runtime. The Go +// port is responsible for the // per-DSL-state wipe (step 2): it transforms the persisted DSL // saved in user_canvas.dsl, the same way the Python handler does // before writing it back via UserCanvasService.update_by_id. diff --git a/internal/agent/runtime/state.go b/internal/agent/runtime/state.go index 70706b26fc..f1f9dc1c03 100644 --- a/internal/agent/runtime/state.go +++ b/internal/agent/runtime/state.go @@ -70,13 +70,13 @@ type CanvasState struct { Globals map[string]any CancelFlag *atomic.Bool RunID string - TaskID string + SessionID string } // NewCanvasState returns a zero-valued CanvasState with all maps allocated. // The atomic CancelFlag is allocated eagerly so nodes can safely poll it // even before any cancel signal has been wired. -func NewCanvasState(runID, taskID string) *CanvasState { +func NewCanvasState(runID, sessionID string) *CanvasState { s := &CanvasState{ activeHistoryIndex: -1, Outputs: make(map[string]map[string]any), @@ -89,7 +89,7 @@ func NewCanvasState(runID, taskID string) *CanvasState { Globals: make(map[string]any), CancelFlag: &atomic.Bool{}, RunID: runID, - TaskID: taskID, + SessionID: sessionID, } s.EnsureSysDate() return s @@ -142,7 +142,7 @@ type canvasStateJSON struct { Globals map[string]any `json:"globals,omitempty"` CancelFlag bool `json:"cancel_flag"` RunID string `json:"run_id"` - TaskID string `json:"task_id"` + SessionID string `json:"session_id"` } // MarshalJSON serialises the CanvasState for eino's StatePre/Post @@ -181,7 +181,7 @@ func (s *CanvasState) MarshalJSON() ([]byte, error) { Globals: s.Globals, CancelFlag: s.CancelFlag != nil && s.CancelFlag.Load(), RunID: s.RunID, - TaskID: s.TaskID, + SessionID: s.SessionID, } // Use SafeJSONMarshal to handle non-serializable values (funcs, // channels) that may have leaked into state maps. Mirrors the @@ -228,7 +228,7 @@ func (s *CanvasState) UnmarshalJSON(b []byte) error { } s.CancelFlag.Store(snap.CancelFlag) s.RunID = snap.RunID - s.TaskID = snap.TaskID + s.SessionID = snap.SessionID return nil } diff --git a/internal/agent/runtime/state_test.go b/internal/agent/runtime/state_test.go index d4896b32a1..ea74e49f76 100644 --- a/internal/agent/runtime/state_test.go +++ b/internal/agent/runtime/state_test.go @@ -26,10 +26,10 @@ import ( // interrupt path's "failed to marshal state: unknown type: // runtime.CanvasState" error). Every field on CanvasState must // round-trip without losing the map values, the CancelFlag bool, -// and the RunID / TaskID strings. +// and the RunID / SessionID strings. func TestCanvasState_MarshalUnmarshalJSON(t *testing.T) { t.Parallel() - src := NewCanvasState("run-1", "task-1") + src := NewCanvasState("run-1", "session-1") src.Sys["query"] = "hello" src.Env["counter"] = 0.0 src.CancelFlag.Store(true) @@ -58,8 +58,8 @@ func TestCanvasState_MarshalUnmarshalJSON(t *testing.T) { if dst.RunID != "run-1" { t.Errorf("RunID = %q, want %q", dst.RunID, "run-1") } - if dst.TaskID != "task-1" { - t.Errorf("TaskID = %q, want %q", dst.TaskID, "task-1") + if dst.SessionID != "session-1" { + t.Errorf("SessionID = %q, want %q", dst.SessionID, "session-1") } if len(dst.Path) != 2 || dst.Path[0] != "begin_0" || dst.Path[1] != "message_0" { t.Errorf("Path = %v, want [begin_0 message_0]", dst.Path) diff --git a/internal/agent/sandbox/manager_client.go b/internal/agent/sandbox/manager_client.go index 8ee99a7fbc..b2a85064c6 100644 --- a/internal/agent/sandbox/manager_client.go +++ b/internal/agent/sandbox/manager_client.go @@ -3,6 +3,7 @@ package sandbox import ( "context" "fmt" + "time" agenttool "ragflow/internal/agent/tool" ) @@ -33,7 +34,11 @@ func (c *ManagerClient) ExecuteCode(ctx context.Context, req agenttool.SandboxRe if err != nil { return nil, err } - defer func() { _ = provider.DestroyInstance(context.Background(), inst) }() + defer func() { + cleanupCtx, cancel := context.WithTimeout(context.WithoutCancel(ctx), 10*time.Second) + defer cancel() + _ = provider.DestroyInstance(cleanupCtx, inst) + }() timeout := req.Timeout if timeout == 0 { diff --git a/internal/agent/sandbox/ssh.go b/internal/agent/sandbox/ssh.go index 21011c54bb..4f5d17e22b 100644 --- a/internal/agent/sandbox/ssh.go +++ b/internal/agent/sandbox/ssh.go @@ -234,11 +234,11 @@ func (p *SSHProvider) CreateInstance(ctx context.Context, template string) (*San remoteBase := p.workDir remoteWorkDir := path.Join(remoteBase, "ragflow-ssh-"+uuid.NewString()) // Create the work_dir and an artifacts/ subdir on the remote. - if err := p.remoteMkdirAll(client, remoteWorkDir); err != nil { + if err := p.remoteMkdirAll(ctx, client, remoteWorkDir); err != nil { _ = client.Close() return nil, fmt.Errorf("ssh: mkdir remote work_dir: %w", err) } - if err := p.remoteMkdirAll(client, path.Join(remoteWorkDir, "artifacts")); err != nil { + if err := p.remoteMkdirAll(ctx, client, path.Join(remoteWorkDir, "artifacts")); err != nil { _ = client.Close() return nil, fmt.Errorf("ssh: mkdir remote artifacts: %w", err) } @@ -321,7 +321,7 @@ func (p *SSHProvider) ExecuteCode( bin = p.nodeBin } remoteScriptPath := path.Join(instance.remoteWorkDir, scriptName) - if err := p.remoteWriteFile(instance.client, remoteScriptPath, wrapped); err != nil { + if err := p.remoteWriteFile(ctx, instance.client, remoteScriptPath, wrapped); err != nil { return nil, fmt.Errorf("ssh: upload script: %w", err) } @@ -348,7 +348,7 @@ func (p *SSHProvider) ExecuteCode( cleanedStdout, structured := ExtractStructuredResult(stdout) // Collect artifacts. - artifacts, err := p.collectArtifacts(instance.client, path.Join(instance.remoteWorkDir, "artifacts")) + artifacts, err := p.collectArtifacts(ctx, instance.client, path.Join(instance.remoteWorkDir, "artifacts")) if err != nil { return nil, fmt.Errorf("ssh: collect artifacts: %w", err) } @@ -413,12 +413,7 @@ func (p *SSHProvider) HealthCheck(ctx context.Context) error { return err } defer client.Close() - sess, err := client.NewSession() - if err != nil { - return fmt.Errorf("ssh: open session: %w", err) - } - defer sess.Close() - if err := sess.Run("true"); err != nil { + if _, _, _, err := p.runRemoteCommand(ctx, client, "true", minTimeout(p.timeout, 10)); err != nil { return fmt.Errorf("ssh: run health probe: %w", err) } return nil @@ -459,11 +454,42 @@ func (p *SSHProvider) dial(ctx context.Context) (*ssh.Client, error) { Timeout: time.Duration(p.timeout) * time.Second, } addr := net.JoinHostPort(p.host, strconv.Itoa(p.port)) - client, err := ssh.Dial("tcp", addr, cfg) + // ssh.Dial has no context-aware variant. Establish the TCP connection + // with DialContext, then run the SSH handshake with a deadline and close + // the socket if the caller cancels while the handshake is in progress. + dialer := net.Dialer{Timeout: time.Duration(p.timeout) * time.Second} + conn, err := dialer.DialContext(ctx, "tcp", addr) if err != nil { return nil, fmt.Errorf("ssh: dial %s: %w", addr, err) } - return client, nil + deadline := time.Now().Add(time.Duration(p.timeout) * time.Second) + if p.timeout > 0 { + _ = conn.SetDeadline(deadline) + } + type handshakeResult struct { + clientConn ssh.Conn + channels <-chan ssh.NewChannel + requests <-chan *ssh.Request + err error + } + handshake := make(chan handshakeResult, 1) + go func() { + clientConn, channels, requests, handshakeErr := ssh.NewClientConn(conn, addr, cfg) + handshake <- handshakeResult{clientConn: clientConn, channels: channels, requests: requests, err: handshakeErr} + }() + var result handshakeResult + select { + case <-ctx.Done(): + _ = conn.Close() + return nil, fmt.Errorf("ssh: dial %s: %w", addr, ctx.Err()) + case result = <-handshake: + } + if result.err != nil { + _ = conn.Close() + return nil, fmt.Errorf("ssh: dial %s: %w", addr, result.err) + } + _ = conn.SetDeadline(time.Time{}) + return ssh.NewClient(result.clientConn, result.channels, result.requests), nil } // hostKeyCallback builds an ssh.HostKeyCallback backed by an OpenSSH @@ -502,7 +528,30 @@ func (p *SSHProvider) runRemoteCommand(ctx context.Context, client *ssh.Client, // from shq()-escaped arguments only (see callers above); user // input never reaches the shell unsanitized. // codeql[go/command-injection] False positive: command is built - if err := sess.Run(command); err != nil { + if err := sess.Start(command); err != nil { + return stdoutBuf.String(), stderrBuf.String(), -1, err + } + runCtx := ctx + cancel := func() {} + if timeoutSec > 0 { + runCtx, cancel = context.WithTimeout(ctx, time.Duration(timeoutSec)*time.Second) + } + defer cancel() + done := make(chan error, 1) + go func() { done <- sess.Wait() }() + + select { + case err = <-done: + case <-runCtx.Done(): + // Closing the SSH session interrupts Wait and asks the remote server to + // terminate the associated command channel. Wait for the goroutine so + // the output builders are no longer being written before returning. + _ = sess.Signal(ssh.SIGKILL) + _ = sess.Close() + <-done + return stdoutBuf.String(), stderrBuf.String(), -1, runCtx.Err() + } + if err != nil { // ssh.ExitError carries the remote exit code; we surface // it as a normal non-zero exit (the caller can branch on // the ExitCode field). @@ -518,8 +567,8 @@ func (p *SSHProvider) runRemoteCommand(ctx context.Context, client *ssh.Client, // remoteMkdirAll runs `mkdir -p` on the remote. The Python // side uses paramiko's mkdir + walk-and-mkdir loop; SSH exec // with `mkdir -p` is simpler and equivalent. -func (p *SSHProvider) remoteMkdirAll(client *ssh.Client, remotePath string) error { - _, stderr, exitCode, err := p.runRemoteCommand(context.Background(), client, +func (p *SSHProvider) remoteMkdirAll(ctx context.Context, client *ssh.Client, remotePath string) error { + _, stderr, exitCode, err := p.runRemoteCommand(ctx, client, fmt.Sprintf("mkdir -p %s", shq(remotePath)), minTimeout(p.timeout, 10), ) @@ -539,13 +588,13 @@ func (p *SSHProvider) remoteMkdirAll(client *ssh.Client, remotePath string) erro // (>1 MiB) this is inefficient vs. SFTP; the threshold is // intentionally not implemented here — Python's paramiko // also writes via SFTP for the same reason. -func (p *SSHProvider) remoteWriteFile(client *ssh.Client, remotePath, content string) error { +func (p *SSHProvider) remoteWriteFile(ctx context.Context, client *ssh.Client, remotePath, content string) error { const tag = "__RAGFLOW_SSH_EOF__" cmd := fmt.Sprintf( "cat > %s <<'%s'\n%s\n%s", shq(remotePath), tag, content, tag, ) - _, stderr, exitCode, err := p.runRemoteCommand(context.Background(), client, cmd, p.timeout) + _, stderr, exitCode, err := p.runRemoteCommand(ctx, client, cmd, p.timeout) if err != nil { return err } @@ -557,8 +606,8 @@ func (p *SSHProvider) remoteWriteFile(client *ssh.Client, remotePath, content st // remoteReadFile reads a remote file's content as a string. // Used by collectArtifacts. -func (p *SSHProvider) remoteReadFile(client *ssh.Client, remotePath string) (string, error) { - stdout, stderr, exitCode, err := p.runRemoteCommand(context.Background(), client, +func (p *SSHProvider) remoteReadFile(ctx context.Context, client *ssh.Client, remotePath string) (string, error) { + stdout, stderr, exitCode, err := p.runRemoteCommand(ctx, client, fmt.Sprintf("cat %s", shq(remotePath)), p.timeout, ) @@ -575,7 +624,7 @@ func (p *SSHProvider) remoteReadFile(client *ssh.Client, remotePath string) (str // is `namesizemode` per line, sorted lexically by the // remote `find` call. We use `find` rather than `ls -la` because // its output is unambiguous across distros (no header rows). -func (p *SSHProvider) remoteListDir(client *ssh.Client, remotePath string) ([]remoteEntry, error) { +func (p *SSHProvider) remoteListDir(ctx context.Context, client *ssh.Client, remotePath string) ([]remoteEntry, error) { // -mindepth 1 / -maxdepth 1: only direct children, not // the dir itself. -printf 'P\t%s\t%m\n' is the GNU find // format; the leading P is a literal path placeholder @@ -587,7 +636,7 @@ func (p *SSHProvider) remoteListDir(client *ssh.Client, remotePath string) ([]re "find %s -mindepth 1 -maxdepth 1 -printf '%%p\\t%%s\\t%%m\\n'", shq(remotePath), ) - stdout, stderr, exitCode, err := p.runRemoteCommand(context.Background(), client, cmd, p.timeout) + stdout, stderr, exitCode, err := p.runRemoteCommand(ctx, client, cmd, p.timeout) if err != nil { return nil, err } @@ -626,8 +675,8 @@ type remoteEntry struct { // collectArtifacts walks the remote artifacts/ dir and returns // the list of files as {name, content_b64, mime_type, size}. // Enforces the same limits the local provider does. -func (p *SSHProvider) collectArtifacts(client *ssh.Client, root string) ([]map[string]any, error) { - entries, err := p.remoteListDir(client, root) +func (p *SSHProvider) collectArtifacts(ctx context.Context, client *ssh.Client, root string) ([]map[string]any, error) { + entries, err := p.remoteListDir(ctx, client, root) if err != nil { return nil, err } @@ -636,7 +685,7 @@ func (p *SSHProvider) collectArtifacts(client *ssh.Client, root string) ([]map[s remote := path.Join(root, e.Name) // Mode bits: S_ISDIR = 0o040000, S_ISREG = 0o100000. if e.Mode&0o170000 == 0o040000 { - sub, err := p.collectArtifacts(client, remote) + sub, err := p.collectArtifacts(ctx, client, remote) if err != nil { return nil, err } @@ -656,7 +705,7 @@ func (p *SSHProvider) collectArtifacts(client *ssh.Client, root string) ([]map[s if _, ok := allowedArtifactExts[ext]; !ok { return nil, fmt.Errorf("unsupported artifact type: %s", e.Name) } - body, err := p.remoteReadFile(client, remote) + body, err := p.remoteReadFile(ctx, client, remote) if err != nil { return nil, err } diff --git a/internal/dao/api_token.go b/internal/dao/api_token.go index cbc2c76d50..eab665a643 100644 --- a/internal/dao/api_token.go +++ b/internal/dao/api_token.go @@ -171,6 +171,20 @@ func (dao *API4ConversationDAO) GetBySessionID(ctx context.Context, db *gorm.DB, return &result, nil } +// GetByID returns a conversation without requiring the caller to know its +// agent. It is used when the session itself is the authorization resource. +func (dao *API4ConversationDAO) GetByID(ctx context.Context, db *gorm.DB, id string) (*entity.API4Conversation, error) { + var result entity.API4Conversation + tx := db.WithContext(ctx).Where("id = ?", id).Find(&result) + if tx.Error != nil { + return nil, tx.Error + } + if tx.RowsAffected == 0 { + return nil, nil + } + return &result, nil +} + // ListIDsByAgentID lists conversation IDs for one agent. func (dao *API4ConversationDAO) ListIDsByAgentID(ctx context.Context, db *gorm.DB, agentID string) ([]string, error) { var ids []string diff --git a/internal/handler/agent.go b/internal/handler/agent.go index 635587d989..08fd8ef27d 100644 --- a/internal/handler/agent.go +++ b/internal/handler/agent.go @@ -37,6 +37,7 @@ import ( "ragflow/internal/entity" "ragflow/internal/service" "ragflow/internal/service/file" + "ragflow/internal/utility" dslpkg "ragflow/internal/agent/dsl" ) @@ -200,8 +201,9 @@ func (h *AgentHandler) ListAgents(c *gin.Context) { // mapAgentError normalises service-layer errors onto the existing // {code, data, message} response envelope used by every other handler. // -// Three classes: +// Four classes: // - service.ErrAgentNotOwner -> "Only the owner..." (DELETE only, 103) +// - service.ErrAgentSessionBusy -> "session already running" (103) // - dao.ErrUserCanvasNotFound -> "Make sure you have permission..." (103) // - service.ErrAgentStorageError -> "Internal storage error" (500) // @@ -219,6 +221,9 @@ func mapAgentError(err error) (common.ErrorCode, string) { if errors.Is(err, service.ErrAgentNotOwner) { return common.CodeOperatingError, "Only the owner of the agent is authorized for this operation." } + if errors.Is(err, service.ErrAgentSessionBusy) { + return common.CodeOperatingError, "This agent session is already running." + } if errors.Is(err, dao.ErrUserCanvasNotFound) || errors.Is(err, dao.ErrUserCanvasVersionNotFound) { return common.CodeOperatingError, "Make sure you have permission to access the agent." @@ -390,6 +395,11 @@ func (h *AgentHandler) RunAgent(c *gin.Context) { canvasID := c.Param("canvas_id") version := c.Query("version") sessionID := c.Query("session_id") + if sessionID == "" { + // Allocate the ordinary-Agent session identity at the HTTP boundary. + // Persistence of the session record remains owned by AgentService.RunAgent. + sessionID = utility.GenerateToken() + } userInput := readUserInput(c) events, err := h.chatRunner.RunAgent(c.Request.Context(), user.ID, canvasID, sessionID, version, userInput, nil) @@ -469,21 +479,24 @@ func sanitiseRunEventError(data string) string { return data } -// CancelAgent signals the in-flight run to stop. -// @Summary Cancel Agent Run +// CancelSessionRun cancels one ordinary Agent run by session id. +// @Summary Cancel Agent Session Run // @Tags agents // @Produce json -// @Param canvas_id path string true "canvas id" +// @Param session_id path string true "session id" // @Success 200 {object} map[string]interface{} -// @Router /api/v1/agents/{canvas_id}/run [delete] -func (h *AgentHandler) CancelAgent(c *gin.Context) { +// @Router /api/v1/tasks/{session_id}/cancel [post] +func (h *AgentHandler) CancelSessionRun(c *gin.Context) { user, code, msg := GetUser(c) if code != common.CodeSuccess { common.ResponseWithCodeData(c, code, nil, msg) return } - canvasID := c.Param("canvas_id") - if err := h.agentService.CancelAgent(c.Request.Context(), user.ID, canvasID); err != nil { + if h.agentService == nil { + common.ResponseWithCodeData(c, common.CodeServerError, nil, "agent service unavailable") + return + } + if err := h.agentService.CancelSessionRun(c.Request.Context(), user.ID, c.Param("session_id")); err != nil { ec, em := mapAgentError(err) common.ResponseWithCodeData(c, ec, nil, em) return @@ -1007,6 +1020,12 @@ func (h *AgentHandler) AgentChatCompletions(c *gin.Context) { zap.String("session_id", req.SessionID), }, userInputMeta(userInput)...)..., ) + if req.SessionID == "" { + // Keep the effective session available to the non-stream response and + // to the task_id=session_id wire alias even when the canvas emits no + // events (for example an empty query). + req.SessionID = utility.GenerateToken() + } events, err := h.chatRunner.RunAgent(c.Request.Context(), user.ID, req.AgentID, req.SessionID, "", userInput, req.Files) if err != nil { @@ -1044,7 +1063,6 @@ func (h *AgentHandler) AgentChatCompletions(c *gin.Context) { zap.String("session_id", req.SessionID), zap.String("event_type", ev.Type), zap.String("message_id", ev.MessageID), - zap.String("task_id", ev.TaskID), ) if err := service.WriteChatbotRunEvent(c.Writer, ev); err != nil { common.Debug("agent chat completions: client disconnected", @@ -1183,7 +1201,7 @@ func (h *AgentHandler) AgentChatCompletions(c *gin.Context) { "data": ansData, "message_id": finalAns.MessageID, "created_at": finalAns.CreatedAt, - "task_id": finalAns.TaskID, + "task_id": finalAns.SessionID, "session_id": finalAns.SessionID, } common.SuccessWithData(c, result, "success") diff --git a/internal/handler/agent_test.go b/internal/handler/agent_test.go index 02be7b0868..e63e7f1bc4 100644 --- a/internal/handler/agent_test.go +++ b/internal/handler/agent_test.go @@ -452,9 +452,6 @@ func (f *fullFakeAgentService) RunAgent(context.Context, string, string, string, close(ch) return ch, nil } -func (f *fullFakeAgentService) CancelAgent(context.Context, string, string) error { - return nil -} func (f *fullFakeAgentService) PublishAgent(context.Context, string, string, *service.PublishAgentRequest) (*entity.UserCanvasVersion, error) { return f.version, nil } @@ -500,7 +497,6 @@ func TestAgentHandler_RoutesRegistered(t *testing.T) { g.PUT("/:canvas_id", func(c *gin.Context) { c.Status(http.StatusOK) }) g.DELETE("/:canvas_id", func(c *gin.Context) { c.Status(http.StatusOK) }) g.POST("/:canvas_id/run", func(c *gin.Context) { c.Status(http.StatusOK) }) - g.DELETE("/:canvas_id/run", func(c *gin.Context) { c.Status(http.StatusOK) }) g.POST("/:canvas_id/publish", func(c *gin.Context) { c.Status(http.StatusOK) }) g.GET("/:canvas_id/versions", func(c *gin.Context) { c.Status(http.StatusOK) }) g.GET("/:canvas_id/versions/:version_id", func(c *gin.Context) { c.Status(http.StatusOK) }) @@ -516,14 +512,13 @@ func TestAgentHandler_RoutesRegistered(t *testing.T) { {http.MethodPut, "/api/v1/agents/abc"}, {http.MethodDelete, "/api/v1/agents/abc"}, {http.MethodPost, "/api/v1/agents/abc/run"}, - {http.MethodDelete, "/api/v1/agents/abc/run"}, {http.MethodPost, "/api/v1/agents/abc/publish"}, {http.MethodGet, "/api/v1/agents/abc/versions"}, {http.MethodGet, "/api/v1/agents/abc/versions/v1"}, {http.MethodDelete, "/api/v1/agents/abc/versions/v1"}, } - if len(routes) != 11 { - t.Fatalf("expected 11 routes, listed %d", len(routes)) + if len(routes) != 10 { + t.Fatalf("expected 10 routes, listed %d", len(routes)) } for _, rt := range routes { w := httptest.NewRecorder() @@ -731,7 +726,8 @@ func (s *stubChatRunner) RunAgent(_ context.Context, _, _, _, _ string, _ any, _ // path: the handler streams canvas.RunEvent frames as // `data: {...}\n\n` with a trailing `data:[DONE]\n\n` terminator. // The frame shape is the Python agent-canvas envelope -// {event,message_id,task_id,session_id,data:{content}}. See +// {event,message_id,task_id,session_id,data:{content}}. task_id is a wire alias +// for session_id. See // service.WriteChatbotRunEvent. // // The stubChatRunner emits one `message` frame and one `done` frame @@ -748,7 +744,7 @@ func TestAgentChatCompletions_StreamSetsContentType(t *testing.T) { c.Set("user_id", "u1") runner := &stubChatRunner{events: []canvas.RunEvent{ - {Type: "message", MessageID: "msg-1", TaskID: "task-1", SessionID: "sess-1", Data: `{"content":"hi back","reference":[]}`}, + {Type: "message", MessageID: "msg-1", SessionID: "sess-1", Data: `{"content":"hi back","reference":[]}`}, {Type: "done", Data: ""}, }} h := &AgentHandler{chatRunner: runner} @@ -760,7 +756,7 @@ func TestAgentChatCompletions_StreamSetsContentType(t *testing.T) { body := w.Body.String() if !strings.Contains(body, `"event":"message"`) || !strings.Contains(body, `"message_id":"msg-1"`) || - !strings.Contains(body, `"task_id":"task-1"`) || + !strings.Contains(body, `"task_id":"sess-1"`) || !strings.Contains(body, `"session_id":"sess-1"`) || !strings.Contains(body, `"content":"hi back"`) { t.Errorf("body should contain flat agent event with content, got %q", body) @@ -781,7 +777,7 @@ func TestAgentChatCompletions_StreamAddsDoneWhenRunnerCloses(t *testing.T) { c.Set("user_id", "u1") runner := &stubChatRunner{events: []canvas.RunEvent{ - {Type: "message", MessageID: "msg-1", TaskID: "task-1", SessionID: "sess-1", Data: `{"content":"hi back"}`}, + {Type: "message", MessageID: "msg-1", SessionID: "sess-1", Data: `{"content":"hi back"}`}, }} h := &AgentHandler{chatRunner: runner} h.AgentChatCompletions(c) @@ -807,7 +803,7 @@ func TestRunAgent_StreamAddsDoneWhenRunnerCloses(t *testing.T) { c.Set("user_id", "u1") runner := &stubChatRunner{events: []canvas.RunEvent{ - {Type: "message", MessageID: "msg-1", TaskID: "task-1", SessionID: "sess-1", Data: `{"content":"hi back"}`}, + {Type: "message", MessageID: "msg-1", SessionID: "sess-1", Data: `{"content":"hi back"}`}, }} h := &AgentHandler{chatRunner: runner} h.RunAgent(c) @@ -837,7 +833,7 @@ func TestAgentChatCompletions_DefaultBranchNonStreaming(t *testing.T) { c.Set("user_id", "u1") runner := &stubChatRunner{events: []canvas.RunEvent{ - {Type: "message", MessageID: "msg-2", TaskID: "task-2", SessionID: "sess-2", Data: `{"content":"hello back","reference":[]}`}, + {Type: "message", MessageID: "msg-2", SessionID: "sess-2", Data: `{"content":"hello back","reference":[]}`}, {Type: "done", Data: ""}, }} h := &AgentHandler{chatRunner: runner} @@ -860,6 +856,44 @@ func TestAgentChatCompletions_DefaultBranchNonStreaming(t *testing.T) { } } +type emptySessionCaptureRunner struct { + sessionID string +} + +func (r *emptySessionCaptureRunner) RunAgent(_ context.Context, _, _, sessionID, _ string, _ any, _ []map[string]interface{}) (<-chan canvas.RunEvent, error) { + r.sessionID = sessionID + ch := make(chan canvas.RunEvent) + close(ch) + return ch, nil +} + +func TestAgentChatCompletions_EmptyOutputReturnsGeneratedSession(t *testing.T) { + gin.SetMode(gin.TestMode) + w := httptest.NewRecorder() + c, _ := gin.CreateTestContext(w) + c.Request = httptest.NewRequest("POST", "/api/v1/agents/chat/completions", + strings.NewReader(`{"agent_id":"a1","query":""}`)) + c.Request.Header.Set("Content-Type", "application/json") + c.Set("user", &entity.User{ID: "u1"}) + c.Set("user_id", "u1") + + runner := &emptySessionCaptureRunner{} + h := &AgentHandler{chatRunner: runner} + h.AgentChatCompletions(c) + + if runner.sessionID == "" { + t.Fatal("handler passed an empty session id to RunAgent") + } + var response map[string]any + if err := json.Unmarshal(w.Body.Bytes(), &response); err != nil { + t.Fatalf("decode response: %v", err) + } + data, _ := response["data"].(map[string]any) + if got, _ := data["session_id"].(string); got != runner.sessionID { + t.Fatalf("empty-output session_id = %q, want %q", got, runner.sessionID) + } +} + // TestAgentChatCompletions_DerivesUserInputFromMessages covers the // fallback path: the request omits `query` but supplies `messages` // with a trailing user message. The handler must use that message's diff --git a/internal/handler/agent_wait_for_user_test.go b/internal/handler/agent_wait_for_user_test.go index 430278aff3..4cb34e73f8 100644 --- a/internal/handler/agent_wait_for_user_test.go +++ b/internal/handler/agent_wait_for_user_test.go @@ -114,7 +114,6 @@ func (f *waitFakeAgentService) RunAgent(ctx context.Context, userID, canvasID, s }), nil } -func (f *waitFakeAgentService) CancelAgent(context.Context, string, string) error { return nil } func (f *waitFakeAgentService) PublishAgent(context.Context, string, string, *service.PublishAgentRequest) (*entity.UserCanvasVersion, error) { return &entity.UserCanvasVersion{}, nil } diff --git a/internal/handler/agent_webhook.go b/internal/handler/agent_webhook.go index 942f33baf2..2173533300 100644 --- a/internal/handler/agent_webhook.go +++ b/internal/handler/agent_webhook.go @@ -68,6 +68,8 @@ import ( "ragflow/internal/agent/canvas" "ragflow/internal/common" rediscli "ragflow/internal/engine/redis" + "ragflow/internal/service" + "ragflow/internal/utility" "strconv" "strings" "time" @@ -224,7 +226,7 @@ func (h *AgentHandler) Webhook(c *gin.Context) { } // Detached background run — does NOT inherit c.Request.Context() // so a client disconnect does not cancel the canvas run. - go h.runWebhookDetached(cv, clean, isTest, startTs) + go h.runWebhookDetached(c.Request.Context(), cv, clean, isTest, startTs) c.Data(status, contentType, payload) return } @@ -494,17 +496,19 @@ func renderImmediatelyResponse(cfg map[string]any) (int, string, []byte, error) return status, "text/plain", []byte(bodyTpl), nil } -// runWebhookDetached runs the canvas in the background. It uses -// context.Background() with a 5-minute timeout (NOT -// c.Request.Context()) so a client disconnect does NOT cancel the run. +// runWebhookDetached runs the canvas with a five-minute timeout. It preserves +// request-scoped context values while intentionally detaching cancellation so +// a client disconnect does not cancel an Immediate webhook run. // Trace events are appended to the redis key when isTest is true. // // Mirrors python: agent_api.py:2123-2175 (the asyncio.create_task body // inside the Immediately branch). func (h *AgentHandler) runWebhookDetached( - cv *entity.UserCanvas, payload map[string]any, isTest bool, startTs time.Time, + parent context.Context, cv *entity.UserCanvas, payload map[string]any, isTest bool, startTs time.Time, ) { - ctx, cancel := context.WithTimeout(context.Background(), 5*time.Minute) + sessionID := utility.GenerateToken() + parent = service.WithAgentSessionID(parent, sessionID) + ctx, cancel := context.WithTimeout(context.WithoutCancel(parent), 5*time.Minute) defer cancel() events, err := h.loader.RunAgentWithWebhook(ctx, cv.UserID, cv.ID, payload) @@ -513,11 +517,14 @@ func (h *AgentHandler) runWebhookDetached( zap.String("canvas", cv.ID), zap.Error(err)) if isTest { - appendWebhookTrace(cv.ID, startTs, canvas.RunEvent{Type: "error", Data: mustJSON(map[string]any{"message": err.Error()})}) + appendWebhookTrace(cv.ID, startTs, canvas.RunEvent{Type: "error", SessionID: sessionID, Data: mustJSON(map[string]any{"message": err.Error()})}) } return } for ev := range events { + if ev.SessionID == "" { + ev.SessionID = sessionID + } if isTest { appendWebhookTrace(cv.ID, startTs, ev) } @@ -538,21 +545,28 @@ func (h *AgentHandler) runWebhookSync( isTest bool, startTs time.Time, ) webhookSyncResult { status := 200 + sessionID := utility.GenerateToken() + ctx = service.WithAgentSessionID(ctx, sessionID) events, err := h.loader.RunAgentWithWebhook(ctx, cv.UserID, cv.ID, payload) if err != nil { if isTest { - appendWebhookTrace(cv.ID, startTs, canvas.RunEvent{Type: "error", Data: mustJSON(map[string]any{"message": err.Error()})}) - appendWebhookTrace(cv.ID, startTs, canvas.RunEvent{Type: "finished", Data: mustJSON(map[string]any{"success": false})}) + appendWebhookTrace(cv.ID, startTs, canvas.RunEvent{Type: "error", SessionID: sessionID, Data: mustJSON(map[string]any{"message": err.Error()})}) + appendWebhookTrace(cv.ID, startTs, canvas.RunEvent{Type: "finished", SessionID: sessionID, Data: mustJSON(map[string]any{"success": false})}) } return webhookSyncResult{status: http.StatusBadRequest, body: gin.H{ - "code": 400, - "message": err.Error(), - "success": false, + "code": 400, + "message": err.Error(), + "success": false, + "task_id": sessionID, + "session_id": sessionID, }} } contents := []string{} for ev := range events { + if ev.SessionID == "" { + ev.SessionID = sessionID + } if isTest { appendWebhookTrace(cv.ID, startTs, ev) } @@ -585,12 +599,14 @@ func (h *AgentHandler) runWebhookSync( } final := strings.Join(contents, "") if isTest { - appendWebhookTrace(cv.ID, startTs, canvas.RunEvent{Type: "finished", Data: mustJSON(map[string]any{"success": true})}) + appendWebhookTrace(cv.ID, startTs, canvas.RunEvent{Type: "finished", SessionID: sessionID, Data: mustJSON(map[string]any{"success": true})}) } return webhookSyncResult{status: status, body: gin.H{ - "message": final, - "success": true, - "code": status, + "message": final, + "success": true, + "code": status, + "task_id": sessionID, + "session_id": sessionID, }} } @@ -647,10 +663,8 @@ func appendWebhookTrace(agentID string, startTs time.Time, ev canvas.RunEvent) { if ev.MessageID != "" { eventRecord["message_id"] = ev.MessageID } - if ev.TaskID != "" { - eventRecord["task_id"] = ev.TaskID - } if ev.SessionID != "" { + eventRecord["task_id"] = ev.SessionID eventRecord["session_id"] = ev.SessionID } entry["events"] = append(events, eventRecord) diff --git a/internal/router/agent_routes.go b/internal/router/agent_routes.go index 43168e9330..2538b5b3b3 100644 --- a/internal/router/agent_routes.go +++ b/internal/router/agent_routes.go @@ -51,7 +51,6 @@ func RegisterAgentRoutes(g *gin.RouterGroup, h *handler.AgentHandler) { g.PUT("/:canvas_id", h.UpdateAgent) g.DELETE("/:canvas_id", h.DeleteAgent) g.POST("/:canvas_id/run", h.RunAgent) - g.DELETE("/:canvas_id/run", h.CancelAgent) g.POST("/:canvas_id/publish", h.PublishAgent) g.PUT("/:canvas_id/tags", h.UpdateAgentTags) g.POST("/:canvas_id/reset", h.ResetAgent) @@ -101,6 +100,16 @@ func RegisterAgentRoutes(g *gin.RouterGroup, h *handler.AgentHandler) { g.POST("/test_db_connection", h.TestDBConnection) } +// RegisterAgentCancelRoutes registers ordinary Agent session cancellation on +// the existing /api/v1/tasks URI. The path is retained for clients, while the +// parameter and all runtime semantics are session-scoped. +func RegisterAgentCancelRoutes(g *gin.RouterGroup, h *handler.AgentHandler) { + if g == nil || h == nil { + return + } + g.POST("/:session_id/cancel", h.CancelSessionRun) +} + // registerAnyMethod mirrors the Python // `@manager.route(path, methods=["POST","GET","PUT","PATCH","DELETE","HEAD"])` // pattern. Gin has no Match() helper, so we register each verb diff --git a/internal/router/agent_routes_test.go b/internal/router/agent_routes_test.go index 218adb2a49..2be990d5b8 100644 --- a/internal/router/agent_routes_test.go +++ b/internal/router/agent_routes_test.go @@ -26,14 +26,14 @@ import ( "ragflow/internal/handler" ) -// TestAgentRoutes_AllElevenRegistered exercises the 11 Phase 5 agent +// TestAgentRoutesRegistered exercises the ordinary Agent route set // endpoints via the public RegisterAgentRoutes helper, proving that the // route table defined in agent_routes.go is actually wired when called // from a real router. This guards against the regression captured in // the post-Phase-7 code review: the helper was defined but never // invoked from Router.Setup, so 10 of the 11 endpoints returned 404 in // production even though the helper "looked correct". -func TestAgentRoutes_AllElevenRegistered(t *testing.T) { +func TestAgentRoutesRegistered(t *testing.T) { eng := gin.New() g := eng.Group("/api/v1/agents") RegisterAgentRoutes(g, &handler.AgentHandler{}) @@ -48,15 +48,14 @@ func TestAgentRoutes_AllElevenRegistered(t *testing.T) { {http.MethodPut, "/api/v1/agents/abc"}, {http.MethodDelete, "/api/v1/agents/abc"}, {http.MethodPost, "/api/v1/agents/abc/run"}, - {http.MethodDelete, "/api/v1/agents/abc/run"}, {http.MethodPost, "/api/v1/agents/abc/publish"}, {http.MethodGet, "/api/v1/agents/abc/versions"}, {http.MethodGet, "/api/v1/agents/abc/versions/v1"}, {http.MethodDelete, "/api/v1/agents/abc/versions/v1"}, {http.MethodPost, "/api/v1/agents/abc/reset"}, } - if len(cases) != 12 { - t.Fatalf("expected 12 routes, listed %d", len(cases)) + if len(cases) != 11 { + t.Fatalf("expected 11 routes, listed %d", len(cases)) } for _, c := range cases { w := httptest.NewRecorder() @@ -79,3 +78,14 @@ func TestAgentRoutes_NilSafety(t *testing.T) { RegisterAgentRoutes(eng.Group("/agents"), nil) // Reaching here without panicking is the assertion. } + +func TestAgentSessionCancelRouteRegistered(t *testing.T) { + eng := gin.New() + RegisterAgentCancelRoutes(eng.Group("/api/v1/tasks"), &handler.AgentHandler{}) + w := httptest.NewRecorder() + req := httptest.NewRequest(http.MethodPost, "/api/v1/tasks/session-1/cancel", nil) + eng.ServeHTTP(w, req) + if w.Code == http.StatusNotFound { + t.Fatal("POST /api/v1/tasks/:session_id/cancel was not registered") + } +} diff --git a/internal/router/router.go b/internal/router/router.go index 375f19c347..408de9cbca 100644 --- a/internal/router/router.go +++ b/internal/router/router.go @@ -580,6 +580,9 @@ func (r *Router) Setup(engine *gin.Engine) { // Agent routes agents := v1.Group("/agents") RegisterAgentRoutes(agents, r.agentHandler) + // Keep the established /tasks URI while using session_id as the + // ordinary Agent run and cancellation identity. + RegisterAgentCancelRoutes(v1.Group("/tasks"), r.agentHandler) // Plugin routes plugin := v1.Group("/plugin") diff --git a/internal/service/agent.go b/internal/service/agent.go index 64e7b6b7a9..79a630488f 100644 --- a/internal/service/agent.go +++ b/internal/service/agent.go @@ -27,6 +27,7 @@ import ( "sort" "strings" "sync" + "sync/atomic" "time" "github.com/cloudwego/eino/compose" @@ -87,7 +88,24 @@ func (s *AgentService) RunAgentWithWebhook( if payload != nil { ctx = context.WithValue(ctx, webhookPayloadKey{}, payload) } - return s.RunAgent(ctx, userID, canvasID, "", "", "", nil) + return s.RunAgent(ctx, userID, canvasID, AgentSessionIDFromContext(ctx), "", "", nil) +} + +type agentSessionIDContextKey struct{} + +// WithAgentSessionID lets an HTTP boundary allocate the session identity while +// keeping session-record persistence inside AgentService.RunAgent. +func WithAgentSessionID(ctx context.Context, sessionID string) context.Context { + return context.WithValue(ctx, agentSessionIDContextKey{}, sessionID) +} + +// AgentSessionIDFromContext returns an HTTP-assigned session identity, if any. +func AgentSessionIDFromContext(ctx context.Context) string { + if ctx == nil { + return "" + } + sessionID, _ := ctx.Value(agentSessionIDContextKey{}).(string) + return sessionID } func emitAgentMessageEvents(emit func(string, string), answer, thinking string, reference any) { @@ -264,6 +282,10 @@ func splitMessageContent(content string) []string { // ErrAgentNotOwner (owner). var ErrAgentNotOwner = errors.New("agent not owned by user") +// ErrAgentSessionBusy is returned when a second request attempts to run the +// same Agent session before the current run reaches a terminal state. +var ErrAgentSessionBusy = errors.New("agent session is already running") + // ErrAgentStorageError is returned by RunAgent when the underlying // version / canvas / tenant DAO surfaces a non-sentinel error (DB // connectivity, schema drift, deadlock, etc.). The handler's @@ -305,13 +327,28 @@ type AgentService struct { // runTracker records per-run lifecycle (Start / MarkSucceeded / // MarkFailed / MarkCancelled) to Redis hash "agent:run:{id}". - runTracker *canvas.RunTracker + runTracker *canvas.RunTracker + runMu sync.Mutex + activeSessions map[string]*activeAgentRun +} - // runMu and runStreams coordinate active canvas run goroutines so that - // CancelAgent can signal a specific canvas. The map is keyed by canvas - // ID; values are channels that close to signal cancellation. - runMu sync.Mutex - runStreams map[string]chan struct{} +type activeAgentRun struct { + userID string + canvasID string + sessionID string + leaseToken string + cancelRun context.CancelFunc + cancelRequested atomic.Bool +} + +func (r *activeAgentRun) requestCancel() { + if r == nil { + return + } + r.cancelRequested.Store(true) + if r.cancelRun != nil { + r.cancelRun() + } } // NewAgentService create agent service @@ -345,7 +382,7 @@ func NewAgentServiceWithOptions( versionDAO: dao.NewUserCanvasVersionDAO(), api4ConversationDAO: dao.NewAPI4ConversationDAO(), runner: canvas.NewRunner(), - runStreams: make(map[string]chan struct{}), + activeSessions: make(map[string]*activeAgentRun), checkpointStore: cp, stateSerializer: ser, runTracker: rt, @@ -699,8 +736,9 @@ func (s *AgentService) UpdateAgent(ctx context.Context, userID, canvasID string, // extra GET. // // Reset does NOT create a new user_canvas_version row. It also does NOT touch -// the in-flight run state of any currently executing canvas session; that is -// owned by the Python task executor and is out of scope for the Go port. +// the in-flight run state of any currently executing canvas session; active +// session cancellation and lease ownership remain the responsibility of the +// run service, not this DSL reset operation. // // Errors propagate the same way as GetAgent: a missing canvas, or a // canvas that the user has no access to, surfaces as @@ -932,6 +970,114 @@ func (s *AgentService) RunAgent(ctx context.Context, userID, canvasID, sessionID if sessionID == "" { sessionID = utility.GenerateToken() } + runID := runIDFor(canvasID, map[string]any{"session_id": sessionID}) + lockToken := utility.GenerateToken() + runCtx, cancelRun := context.WithCancel(ctx) + active := &activeAgentRun{ + userID: userID, + canvasID: canvasID, + sessionID: sessionID, + leaseToken: lockToken, + cancelRun: cancelRun, + } + releaseLocal := func() { + s.runMu.Lock() + if s.activeSessions[sessionID] == active { + delete(s.activeSessions, sessionID) + } + s.runMu.Unlock() + } + // Make the distributed lease the first run-lifecycle mutation after canvas + // access is authorized. All version, session, and DSL initialization happens + // only after other instances can observe and cancel this starting run. + if s.runTracker != nil { + registered, registerErr := s.runTracker.RegisterActiveSession(ctx, canvas.ActiveSession{ + SessionID: sessionID, + Token: lockToken, + UserID: userID, + CanvasID: canvasID, + RunID: runID, + }) + if registerErr != nil { + cleanupCtx, cleanupCancel := context.WithTimeout(context.WithoutCancel(ctx), time.Second) + _, _ = s.runTracker.ReleaseActiveSession(cleanupCtx, sessionID, lockToken) + cleanupCancel() + releaseLocal() + cancelRun() + return nil, fmt.Errorf("RunAgent: register active session: %w: %w", registerErr, ErrAgentStorageError) + } + if !registered { + releaseLocal() + cancelRun() + return nil, ErrAgentSessionBusy + } + } + + s.runMu.Lock() + if _, exists := s.activeSessions[sessionID]; exists { + s.runMu.Unlock() + if s.runTracker != nil { + cleanupCtx, cleanupCancel := context.WithTimeout(context.WithoutCancel(ctx), time.Second) + _, _ = s.runTracker.ReleaseActiveSession(cleanupCtx, sessionID, lockToken) + cleanupCancel() + } + cancelRun() + return nil, ErrAgentSessionBusy + } + s.activeSessions[sessionID] = active + s.runMu.Unlock() + + if s.runTracker == nil { + // Without the distributed registry, clear a marker left by a prior + // local run before starting the watcher. A concurrent local Cancel also + // calls cancelRun directly, so this cleanup cannot lose that signal. + clearCtx, cancelClear := context.WithTimeout(ctx, time.Second) + _ = canvas.ClearCancel(clearCtx, sessionID) + cancelClear() + } + + releaseRegistration := func() { + cleanupCtx, cleanupCancel := context.WithTimeout(context.WithoutCancel(ctx), time.Second) + if s.runTracker != nil { + _, releaseErr := s.runTracker.ReleaseActiveSession(cleanupCtx, sessionID, lockToken) + if releaseErr != nil { + common.Warn("agent run: release active session failed", zap.String("session_id", sessionID), zap.Error(releaseErr)) + } + } else { + _ = canvas.ClearCancel(cleanupCtx, sessionID) + } + cleanupCancel() + releaseLocal() + } + + // Start cancellation and lease watchers as soon as the active registration + // is visible. This keeps the lease alive during initialization and preserves + // a cancel marker published before Compile/Invoke begins. + watchCtx, cancelWatch := context.WithCancel(context.WithoutCancel(ctx)) + go canvas.WatchCancel(watchCtx, sessionID, active.requestCancel) + if s.runTracker != nil { + go s.runTracker.WatchActiveSession(watchCtx, sessionID, lockToken, active.requestCancel) + } + checkCtx, cancelCheck := context.WithTimeout(context.WithoutCancel(ctx), time.Second) + if requested, checkErr := canvas.CancelRequested(checkCtx, sessionID); checkErr != nil { + common.Warn("agent run: initial cancel check failed", + zap.String("session_id", sessionID), zap.Error(checkErr)) + } else if requested { + active.requestCancel() + } + cancelCheck() + + // Until the Runner goroutine takes ownership, every initialization error + // must release the active lease and any marker written for this token. + registrationHandedOff := false + defer func() { + if registrationHandedOff { + return + } + cancelWatch() + cancelRun() + releaseRegistration() + }() // Load the version row up front so the run is bound to a real DSL. // @@ -1019,6 +1165,7 @@ func (s *AgentService) RunAgent(ctx context.Context, userID, canvasID, sessionID if dsl == nil { dsl = normalisedDSLForRun(versionRow) } + sessionFound := false if sessionID != "" && s.api4ConversationDAO != nil { session, sessionErr := s.api4ConversationDAO.GetBySessionID(ctx, dao.DB, sessionID, canvasID) if sessionErr != nil { @@ -1027,11 +1174,16 @@ func (s *AgentService) RunAgent(ctx context.Context, userID, canvasID, sessionID if session != nil && session.UserID != userID { return nil, fmt.Errorf("RunAgent: session %q not found: %w", sessionID, dao.ErrUserCanvasNotFound) } + sessionFound = session != nil if session != nil && len(session.DSL) > 0 { dsl = dslpkg.NormalizeForRun(session.DSL) } } - if newSession && len(dsl) > 0 { + // A handler may allocate the session id before calling RunAgent so the + // effective id is available even when the run emits no events. Treat an + // absent conversation row as a first touch regardless of who generated the + // id; there is still only one business identity (session_id). + if !sessionFound || newSession { if err = s.createAgentRunSession(ctx, sessionID, userID, canvasID, dsl, versionRow, userInput); err != nil { return nil, fmt.Errorf("RunAgent: create session %q: %w: %w", sessionID, err, ErrAgentStorageError) } @@ -1045,6 +1197,17 @@ func (s *AgentService) RunAgent(ctx context.Context, userID, canvasID, sessionID "session_id": sessionID, "user_id": userID, } + // The stable run id is derived from the canvas and session. It is only a + // checkpoint/status storage key; session_id remains the public run and + // cancellation identity. + // Recover a pending UserFillUp interrupt from Redis when this request + // lands on a different process. + if userInput != nil && s.runTracker != nil { + if interruptID, ok, ierr := s.runTracker.GetInterruptID(ctx, runID); ierr == nil && ok { + root["__resume_interrupt_id__"] = interruptID + root["__resume_data__"] = userInput + } + } if userInput != nil { root["user_input"] = userInput } @@ -1092,7 +1255,59 @@ func (s *AgentService) RunAgent(ctx context.Context, userID, canvasID, sessionID zap.Any("tenantID", root["tenant_id"]), zap.Any("userInput", root["user_input"])) - return s.runner.Run(ctx, run, canvasID, sessionID, userInput, root), nil + // Cancellation of the HTTP request must reach the workflow, but it must + // not stop the Redis watchers while the real Runner goroutine is still + // unwinding. Otherwise a non-cooperative external call can outlive the + // lease, allowing a second process to acquire the same session. The + // detached watcher context is cancelled only after inner has closed and + // cleanup has taken ownership of the lease release. + lifecycleDone := make(chan struct{}) + go func() { + select { + case <-ctx.Done(): + active.requestCancel() + case <-lifecycleDone: + } + }() + inner := s.runner.Run(runCtx, run, canvasID, sessionID, userInput, root) + + out := make(chan canvas.RunEvent, 8) + go func() { + defer close(out) + defer func() { + cancelWatch() + close(lifecycleDone) + if ctx.Err() != nil { + active.cancelRequested.Store(true) + } + cancelRun() + if active.cancelRequested.Load() { + if s.runTracker != nil { + statusCtx, cancelStatus := context.WithTimeout(context.WithoutCancel(ctx), time.Second) + err := s.runTracker.MarkCancelled(statusCtx, runID) + cancelStatus() + if err != nil { + common.Warn("agent run: mark session cancelled failed", + zap.String("session_id", sessionID), zap.Error(err)) + } + } + } + releaseRegistration() + }() + for ev := range inner { + if ev.SessionID == "" { + ev.SessionID = sessionID + } + select { + case out <- ev: + case <-ctx.Done(): + // Keep draining the Runner so its managed workflow can unwind, + // but stop forwarding frames to a disconnected client. + } + } + }() + registrationHandedOff = true + return out, nil } // buildRunFunc assembles the per-run RunFunc the orchestrator (canvas.Runner) @@ -1108,7 +1323,7 @@ func (s *AgentService) RunAgent(ctx context.Context, userID, canvasID, sessionID // answer-extraction contract. // // Nil-versionRow guard: the RunAgent call site treats "no version -// published" as a legal state and passes nil. We extract taskID +// published" as a legal state and passes nil. We extract sessionID // safely and, when both versionRow and dsl are empty, fall back to // a graceful "no published version" placeholder so the SSE surface // still flows (TestRunAgent_NoVersionPublishedPlaceholder pins this @@ -1139,7 +1354,6 @@ func (s *AgentService) buildRunFunc(canvasID string, versionRow *entity.UserCanv // Extract the event channel + metadata injected by Runner.Run. events, _ := root["__events__"].(chan canvas.RunEvent) messageID, _ := root["__message_id__"].(string) - taskID, _ := root["__task_id__"].(string) sessionID, _ := root["__session_id__"].(string) userID, _ := root["user_id"].(string) @@ -1160,7 +1374,6 @@ func (s *AgentService) buildRunFunc(canvasID string, versionRow *entity.UserCanv Type: typ, Data: data, MessageID: messageID, CreatedAt: time.Now().Unix(), - TaskID: taskID, SessionID: sessionID, }) } @@ -1184,10 +1397,6 @@ func (s *AgentService) buildRunFunc(canvasID string, versionRow *entity.UserCanv startedAt := float64(time.Now().UnixNano()) / 1e9 userInput := root["user_input"] - userInputText := "" - if v, ok := userInput.(string); ok { - userInputText = v - } resumeID, isResume := root["__resume_interrupt_id__"].(string) if !isResume || resumeID == "" { @@ -1196,7 +1405,7 @@ func (s *AgentService) buildRunFunc(canvasID string, versionRow *entity.UserCanv } runID := runIDFor(canvasID, root) - state := canvas.NewCanvasState(runID, taskID) + state := canvas.NewCanvasState(runID, sessionID) // Graceful placeholder: no version published AND no DSL. if versionRow == nil && len(dsl) == 0 { @@ -1239,7 +1448,6 @@ func (s *AgentService) buildRunFunc(canvasID string, versionRow *entity.UserCanv ctx2 := canvas.WithRunMeta(ctx, &canvas.RunMeta{ Events: events, MessageID: messageID, - TaskID: taskID, SessionID: sessionID, }) ctx2 = runtime.WithDeferredNodeRegistry(ctx2) @@ -1282,6 +1490,7 @@ func (s *AgentService) buildRunFunc(canvasID string, versionRow *entity.UserCanv if uid, ok := root["user_id"].(string); ok && uid != "" { state.Sys["user_id"] = uid } + state.Sys["agent_id"] = canvasID if tid, ok := root["tenant_id"].(string); ok && tid != "" { state.Sys["tenant_id"] = tid } @@ -1323,7 +1532,7 @@ func (s *AgentService) buildRunFunc(canvasID string, versionRow *entity.UserCanv if s.runTracker != nil { _ = s.runTracker.Start(ctx2, runID, canvasID, - tenantIDFromRoot(root), userInputText) + tenantIDFromRoot(root), "") } // Compile. @@ -1345,7 +1554,6 @@ func (s *AgentService) buildRunFunc(canvasID string, versionRow *entity.UserCanv common.Debug("RunAgent compile err", zap.String("canvas", canvasID), zap.String("session", sessionID), - zap.String("task", taskID), zap.String("run", runID), zap.String("type", fmt.Sprintf("%T", err)), zap.Error(err)) @@ -1381,6 +1589,12 @@ func (s *AgentService) buildRunFunc(canvasID string, versionRow *entity.UserCanv } workflowOutput, invokeErr := cc.Workflow.Invoke(ctx2, map[string]any{"query": wfInput}, invokeOpts...) err = invokeErr + if errors.Is(err, context.Canceled) || errors.Is(ctx2.Err(), context.Canceled) { + // A user stop or client disconnect must not be turned into a + // failed/succeeded run and must not append a synthetic assistant + // message after the workflow has observed cancellation. + return nil, context.Canceled + } if cpID != "" && s.runTracker != nil { _ = s.runTracker.AttachCheckpoint(ctx2, runID, cpID) @@ -1430,12 +1644,15 @@ func (s *AgentService) buildRunFunc(canvasID string, versionRow *entity.UserCanv common.Debug("RunAgent invoke err", zap.String("canvas", canvasID), zap.String("session", sessionID), - zap.String("task", taskID), zap.String("run", runID), zap.String("type", fmt.Sprintf("%T", err)), zap.Error(err)) if canvas.IsInterruptError(err) { - s.markRunFailed(ctx2, runID, "interrupt: "+err.Error()) + resumeID := canvas.RootInterruptID(canvas.ExtractInterruptContexts(err)) + if resumeID != "" && s.runTracker != nil { + _ = s.runTracker.AttachInterrupt(ctx2, runID, resumeID) + _ = s.runTracker.MarkWaiting(ctx2, runID) + } if answer != "" { appendAssistantHistory(state, partialAssistantOutput(answer, downloads)) } @@ -1926,6 +2143,7 @@ func (s *AgentService) markRunSucceeded(ctx context.Context, runID string) { zap.String("run_id", runID), zap.Error(err)) } + _ = s.runTracker.ClearInterruptID(ctx, runID) } // markRunFailed records the run as failed (with reason) via the @@ -1954,24 +2172,54 @@ func normalisedDSLForRun(v *entity.UserCanvasVersion) map[string]any { return dslpkg.NormalizeForRun(v.DSL) } -// CancelAgent signals the in-flight run (if any) for the given canvas to -// stop. It is a no-op when no run is currently registered, or when the -// requesting user has no read access to the canvas. -func (s *AgentService) CancelAgent(ctx context.Context, userID, canvasID string) error { - if _, err := s.loadCanvasForUser(ctx, userID, canvasID); err != nil { - return err - } +// CancelSessionRun is the single ordinary-Agent cancellation path. It cancels +// the local run context when present and publishes a session-scoped Redis +// marker for an owning instance. Unknown and finished sessions are idempotent +// successes. +func (s *AgentService) CancelSessionRun(ctx context.Context, userID, sessionID string) error { s.runMu.Lock() - cancel, ok := s.runStreams[canvasID] + active := s.activeSessions[sessionID] s.runMu.Unlock() - if !ok { + if active != nil { + if active.userID != userID { + return ErrAgentNotOwner + } + active.requestCancel() + if s.runTracker != nil { + if _, err := s.runTracker.RequestCancelActiveSession(ctx, sessionID, active.leaseToken); err != nil { + common.Warn("agent cancel: redis publish failed", zap.String("session_id", sessionID), zap.Error(err)) + } + return nil + } + if err := canvas.RequestCancel(ctx, sessionID); err != nil { + // Local cancellation is already effective; Redis failure only + // prevents cross-instance propagation. + common.Warn("agent cancel: redis publish failed", zap.String("session_id", sessionID), zap.Error(err)) + } return nil } - select { - case <-cancel: - // already closed - default: - close(cancel) + if s.runTracker != nil { + remote, err := s.runTracker.GetActiveSession(ctx, sessionID) + if err != nil { + return fmt.Errorf("agent cancel: read active session: %w: %w", err, ErrAgentStorageError) + } + if err == nil && remote != nil { + if remote.UserID != "" && remote.UserID != userID { + return ErrAgentNotOwner + } + requested, err := s.runTracker.RequestCancelActiveSession(ctx, sessionID, remote.Token) + if err != nil { + return fmt.Errorf("agent cancel: publish remote marker: %w: %w", err, ErrAgentStorageError) + } + if !requested { + return nil + } + return nil + } } + // A persisted conversation is not evidence of an active run. Once both the + // local registration and the distributed lease are absent, cancellation is + // an idempotent no-op. Publishing a session-scoped marker here would race + // completed-run cleanup and could cancel a later run that reuses sessionID. return nil } diff --git a/internal/service/agent_cancel_test.go b/internal/service/agent_cancel_test.go new file mode 100644 index 0000000000..ba75120381 --- /dev/null +++ b/internal/service/agent_cancel_test.go @@ -0,0 +1,174 @@ +// Copyright 2026 The InfiniFlow Authors. All Rights Reserved. +// +// Licensed under the Apache License, Version 2.0 (the "License"); + +package service + +import ( + "encoding/json" + "errors" + "sync/atomic" + "testing" + "time" + + "github.com/alicebob/miniredis/v2" + "github.com/redis/go-redis/v9" + + "ragflow/internal/agent/canvas" + "ragflow/internal/entity" +) + +func newAgentCancelTracker(t *testing.T) (*canvas.RunTracker, *miniredis.Miniredis) { + t.Helper() + mr, err := miniredis.Run() + if err != nil { + t.Fatalf("miniredis.Run: %v", err) + } + t.Cleanup(mr.Close) + client := redis.NewClient(&redis.Options{Addr: mr.Addr()}) + t.Cleanup(func() { _ = client.Close() }) + return canvas.NewRunTrackerWithClient(client, time.Hour), mr +} + +func TestCancelSessionRunLocalPermissionAndIdempotency(t *testing.T) { + svc := NewAgentServiceWithOptions(nil, nil, nil) + active := &activeAgentRun{sessionID: "session-1", userID: "user-a"} + var calls atomic.Int32 + active.cancelRun = func() { + calls.Add(1) + } + svc.activeSessions[active.sessionID] = active + + if err := svc.CancelSessionRun(t.Context(), "user-b", "session-1"); !errors.Is(err, ErrAgentNotOwner) { + t.Fatalf("other user cancel error = %v, want ErrAgentNotOwner", err) + } + if calls.Load() != 0 { + t.Fatal("unauthorized cancel invoked the active cancel func") + } + if err := svc.CancelSessionRun(t.Context(), "user-a", "session-1"); err != nil { + t.Fatalf("owner CancelSessionRun: %v", err) + } + if calls.Load() != 1 || !active.cancelRequested.Load() { + t.Fatalf("local cancel calls=%d requested=%v", calls.Load(), active.cancelRequested.Load()) + } + + delete(svc.activeSessions, "session-1") + if err := svc.CancelSessionRun(t.Context(), "user-a", "session-1"); err != nil { + t.Fatalf("finished session cancel must be idempotent: %v", err) + } + if err := svc.CancelSessionRun(t.Context(), "user-a", "unknown"); err != nil { + t.Fatalf("unknown session cancel must be idempotent: %v", err) + } +} + +func TestCancelSessionRunDoesNotAffectAnotherSession(t *testing.T) { + svc := NewAgentServiceWithOptions(nil, nil, nil) + var callsA, callsB atomic.Int32 + svc.activeSessions["session-a"] = &activeAgentRun{ + userID: "user-a", sessionID: "session-a", cancelRun: func() { callsA.Add(1) }, + } + svc.activeSessions["session-b"] = &activeAgentRun{ + userID: "user-a", sessionID: "session-b", cancelRun: func() { callsB.Add(1) }, + } + if err := svc.CancelSessionRun(t.Context(), "user-a", "session-a"); err != nil { + t.Fatalf("CancelSessionRun: %v", err) + } + if callsA.Load() != 1 || callsB.Load() != 0 { + t.Fatalf("cancel calls A=%d B=%d; want 1, 0", callsA.Load(), callsB.Load()) + } +} + +func TestCancelSessionRunFinishedSessionDoesNotCreateCancelMarker(t *testing.T) { + testDB := setupServiceTestDB(t) + pushServiceDB(t, testDB) + if err := testDB.Create(&entity.API4Conversation{ + ID: "session-persisted", + DialogID: "agent-1", + UserID: "user-a", + Message: json.RawMessage(`[]`), + Reference: json.RawMessage(`[]`), + }).Error; err != nil { + t.Fatalf("create conversation: %v", err) + } + + tracker, mr := newAgentCancelTracker(t) + + ctx := t.Context() + const ( + sessionID = "session-persisted" + runID = "agent-1-session-persisted" + token = "run-owner" + ) + registered, err := tracker.RegisterActiveSession(ctx, canvas.ActiveSession{ + SessionID: sessionID, + Token: token, + UserID: "user-a", + CanvasID: "agent-1", + RunID: runID, + }) + if err != nil || !registered { + t.Fatalf("RegisterActiveSession = %v, %v; want true, nil", registered, err) + } + if err := tracker.Start(ctx, runID, "agent-1", "user-a", ""); err != nil { + t.Fatalf("Start: %v", err) + } + requested, err := tracker.RequestCancelActiveSession(ctx, sessionID, token) + if err != nil || !requested { + t.Fatalf("RequestCancelActiveSession = %v, %v; want true, nil", requested, err) + } + if err := tracker.MarkCancelled(ctx, runID); err != nil { + t.Fatalf("MarkCancelled: %v", err) + } + released, err := tracker.ReleaseActiveSession(ctx, sessionID, token) + if err != nil || !released { + t.Fatalf("ReleaseActiveSession = %v, %v; want true, nil", released, err) + } + + svc := NewAgentServiceWithOptions(nil, nil, tracker) + if err := svc.CancelSessionRun(t.Context(), "user-a", "session-persisted"); err != nil { + t.Fatalf("late owner cancel error = %v", err) + } + if mr.Exists(sessionID + "-cancel") { + t.Fatal("late cancel recreated a marker after the run finished") + } + if err := svc.CancelSessionRun(t.Context(), "user-a", "missing-session"); err != nil { + t.Fatalf("missing session cancel must be idempotent: %v", err) + } + if mr.Exists("missing-session-cancel") { + t.Fatal("unknown-session cancel created a marker") + } +} + +func TestRunAgentInitializationFailureReleasesActiveSession(t *testing.T) { + testDB := setupServiceTestDB(t) + pushServiceDB(t, testDB) + if err := testDB.AutoMigrate(&entity.UserCanvas{}, &entity.UserCanvasVersion{}); err != nil { + t.Fatalf("migrate agent tables: %v", err) + } + if err := testDB.Create(&entity.UserCanvas{ID: "agent-1", UserID: "user-a"}).Error; err != nil { + t.Fatalf("create canvas: %v", err) + } + tracker, mr := newAgentCancelTracker(t) + svc := NewAgentServiceWithOptions(nil, nil, tracker) + + const sessionID = "session-initialization-failure" + if _, err := svc.RunAgent(t.Context(), "user-a", "agent-1", sessionID, "missing-version", "", nil); err == nil { + t.Fatal("RunAgent with missing version returned nil error") + } + active, err := tracker.GetActiveSession(t.Context(), sessionID) + if err != nil { + t.Fatalf("GetActiveSession: %v", err) + } + if active != nil { + t.Fatalf("active session remains after initialization failure: %#v", active) + } + if mr.Exists(sessionID + "-cancel") { + t.Fatal("cancel marker remains after initialization failure") + } + svc.runMu.Lock() + _, locallyActive := svc.activeSessions[sessionID] + svc.runMu.Unlock() + if locallyActive { + t.Fatal("local active session remains after initialization failure") + } +} diff --git a/internal/service/agent_test.go b/internal/service/agent_test.go index 5240d668a7..c442539bc7 100644 --- a/internal/service/agent_test.go +++ b/internal/service/agent_test.go @@ -767,7 +767,7 @@ func TestRunAgent_StorageErrorFromCanvasAccess(t *testing.T) { // contract at the loadCanvasForUser level (not just through // RunAgent). loadCanvasForUser is shared by GetAgent, UpdateAgent, // DeleteAgent, PublishAgent, ListVersions, GetVersion, and -// CancelAgent — sanitising its DAO errors closes the leak in all +// Session cancellation — sanitising its DAO errors closes the leak in all // eight call sites. func TestLoadCanvasForUser_StorageErrorWrap(t *testing.T) { testDB := setupServiceTestDB(t) diff --git a/internal/service/bot_completion.go b/internal/service/bot_completion.go index 48d7713eb1..57243a1134 100644 --- a/internal/service/bot_completion.go +++ b/internal/service/bot_completion.go @@ -181,8 +181,8 @@ func WriteDoneFrame(w http.ResponseWriter) error { // WriteChatbotRunEvent translates one canvas.RunEvent into the flat // Python agent-canvas SSE envelope: // -// data: {"event":"message","message_id":"...","task_id":"...", -// "session_id":"...","created_at":123,"data":{"content":"..."}}\n\n +// data: {"event":"message","message_id":"...","task_id":"session-id", +// "session_id":"session-id","created_at":123,"data":{"content":"..."}}\n\n // // This is intentionally different from WriteChatbotFrame's legacy // chatbot `{code,data:{answer:"..."}}` shape. The agent React page's @@ -228,6 +228,13 @@ func WriteChatbotRunEvent(w http.ResponseWriter, ev canvas.RunEvent) error { "message": msg, "data": false, } + // Keep the error envelope wire-compatible while still correlating the + // failed run. task_id is only the legacy alias; both values are the + // same session identity used by the Go runtime and cancel endpoint. + if ev.SessionID != "" { + payload["task_id"] = ev.SessionID + payload["session_id"] = ev.SessionID + } return writeSSEJSON(w, payload) } @@ -241,10 +248,11 @@ func WriteChatbotRunEvent(w http.ResponseWriter, ev canvas.RunEvent) error { if ev.MessageID != "" { payload["message_id"] = ev.MessageID } - if ev.TaskID != "" { - payload["task_id"] = ev.TaskID - } if ev.SessionID != "" { + // task_id is retained only as a wire-compatible alias for existing Go + // Agent clients. It carries session_id and has no independent runtime, + // Redis, or cancellation identity. + payload["task_id"] = ev.SessionID payload["session_id"] = ev.SessionID } return writeSSEJSON(w, payload) diff --git a/internal/service/bot_completion_history_test.go b/internal/service/bot_completion_history_test.go index 2d41b64713..d3bb217d5f 100644 --- a/internal/service/bot_completion_history_test.go +++ b/internal/service/bot_completion_history_test.go @@ -270,7 +270,6 @@ func TestWriteChatbotRunEvent_UserInputsEvent(t *testing.T) { if err := WriteChatbotRunEvent(rec, canvas.RunEvent{ Type: "user_inputs", MessageID: "msg-1", - TaskID: "task-1", Data: `{"components":[{"id":"email","type":"text","required":true}]}`, SessionID: "sess-1", }); err != nil { @@ -283,8 +282,8 @@ func TestWriteChatbotRunEvent_UserInputsEvent(t *testing.T) { if !strings.Contains(body, `"message_id":"msg-1"`) { t.Errorf("body missing message_id: %s", body) } - if !strings.Contains(body, `"task_id":"task-1"`) { - t.Errorf("body missing task_id: %s", body) + if !strings.Contains(body, `"task_id":"sess-1"`) { + t.Errorf("body missing session-backed task_id alias: %s", body) } if !strings.Contains(body, `"session_id":"sess-1"`) { t.Errorf("body missing session_id: %s", body) @@ -342,6 +341,22 @@ func TestWriteChatbotRunEvent_MessageEventCarriesEvent(t *testing.T) { } } +func TestWriteChatbotRunEvent_ErrorCarriesSessionAlias(t *testing.T) { + rec := &recordingResponseWriter{header: http.Header{}} + if err := WriteChatbotRunEvent(rec, canvas.RunEvent{ + Type: "error", + Data: `{"message":"failed"}`, + SessionID: "sess-error", + }); err != nil { + t.Fatalf("WriteChatbotRunEvent: %v", err) + } + body := rec.body.String() + if !strings.Contains(body, `"task_id":"sess-error"`) || + !strings.Contains(body, `"session_id":"sess-error"`) { + t.Fatalf("error frame missing session-backed aliases: %s", body) + } +} + // TestBotService_ChatbotCompletion_NewSessionSkipsLLM locks in the // share-page handshake behaviour: the front-end opens a shared chat // with an empty question and no session_id only to obtain a session