Skip to content

MUL-4976: feat(daemon): add drain-aware restart for active tasks#5631

Open
wjueyao wants to merge 4 commits into
multica-ai:mainfrom
wjueyao:codex/feat-daemon-drain-restart
Open

MUL-4976: feat(daemon): add drain-aware restart for active tasks#5631
wjueyao wants to merge 4 commits into
multica-ai:mainfrom
wjueyao:codex/feat-daemon-drain-restart

Conversation

@wjueyao

@wjueyao wjueyao commented Jul 19, 2026

Copy link
Copy Markdown
Contributor

What does this PR do?

Adds an opt-in multica daemon restart --drain mode that stops claiming new tasks, lets accepted agent work finish naturally, and only then restarts the daemon.

Thinking path

  1. The current /shutdown path cancels the daemon root context, which is also passed to active agent tasks. The subsequent 30-second wait is cleanup after cancellation, not a real drain.
  2. Server-side runtime_recovery is an important crash fallback, but it still records a failed attempt, redispatches work, and can spend more model tokens resuming it. It should not be the normal restart protocol.
  3. Desktop PR feat(desktop): restart local daemon when bundled CLI version differs #1041 already established the safety direction of deferring a version-mismatch restart while active_task_count > 0; explicit CLI restarts lacked the equivalent capability.
  4. This PR reuses the existing claim barrier, but manual drain acquires it before idle and waits for both claimsInFlight and activeTasks to reach zero. The batch poller increments activeTasks before releasing claimsInFlight, so there is no false-idle window during claim handoff.
  5. Drain uses a dedicated /shutdown/drain endpoint. A new CLI talking to an older daemon therefore gets a safe 404 instead of having an old /shutdown handler ignore a query parameter and stop immediately.
  6. The CLI still owns the final start. This preserves the invoking binary, profile, foreground mode, and restart-time overrides without letting an unauthenticated localhost endpoint choose an executable.

Plain multica daemon restart remains unchanged for backwards compatibility.

Related Issue

Closes #5630

Related context: #1041; orthogonal to the cross-profile ownership work in #5494.

Type of Change

  • Bug fix (non-breaking change that fixes an issue)
  • New feature (non-breaking change that adds functionality)
  • Refactor / code improvement (no behavior change)
  • Documentation update
  • Tests (adding or improving test coverage)
  • CI / infrastructure

Changes Made

  • Add the daemon restart --drain CLI flag and a caller-cancellable, unbounded local request suitable for long-running agent tasks.
  • Add a dedicated drain-shutdown handler and expose draining in local health output.
  • Generalize the claim barrier ownership rules so drain and auto-update cannot steal each other's pause.
  • Cover active-task waiting, claim-to-active handoff, caller cancellation, concurrent barrier conflict, endpoint selection, and immediate-shutdown regression.
  • Document the safe restart option and its design constraints.

How to Test

Automated verification run locally on macOS arm64 with Go 1.26.1:

go test ./internal/daemon ./cmd/multica -count=1
go test -race ./internal/daemon ./cmd/multica -run 'Test(ShutdownHandler|TryEnterClaim_RespectsBarrier|HealthHandlerReportsCLIVersionAndActiveTaskCount|DaemonRestartCommandExposesDrainFlag|RequestDaemonDrainShutdown)' -count=1
go vet ./internal/daemon ./cmd/multica
gofmt -l cmd/multica/cmd_daemon.go cmd/multica/cmd_daemon_test.go internal/daemon/daemon.go internal/daemon/health.go internal/daemon/health_test.go internal/daemon/auto_update_test.go

All commands above pass; gofmt -l produces no output.

Manual reviewer flow:

  1. Start a daemon and assign a task that runs long enough to observe.
  2. Run multica daemon restart --drain in another terminal.
  3. Confirm health reports draining: true, no new task is claimed, and the active task completes without a runtime_recovery failure.
  4. Confirm the old daemon exits and the invoking CLI starts the replacement.
  5. Repeat, but interrupt the waiting CLI; confirm the daemon remains online and resumes claiming.

Risks and compatibility

  • The drain request intentionally has no fixed timeout; the operator can interrupt it, which releases the barrier and resumes claims.
  • A second drain or an overlapping auto-update barrier returns 409 rather than changing barrier ownership.
  • Older daemons do not have /shutdown/drain; they fail safely with 404 and are left running.
  • Default restart behavior and server-side recovery are unchanged.

Checklist

  • I have included a thinking path that traces from project context to this change
  • I have run tests locally and they pass
  • I have added or updated tests where applicable
  • If this change affects the UI, I have included before/after screenshots
  • I have updated relevant documentation to reflect my changes
  • If I added a new runtime / coding tool / UI tab, I synced the change to landing copy and relevant docs
  • If this PR touches Chinese product copy, I checked it against the Chinese copy conventions
  • I have considered and documented risks above
  • I will address all reviewer comments before requesting merge

AI Disclosure

AI tool used: OpenAI Codex

Prompt / approach: Reproduced the daemon lifecycle semantics from main, compared the existing Desktop deferred-restart and auto-update claim-barrier designs, reviewed adjacent PR #5494 for overlap, then implemented test-first. The final diff received separate race/behavior, compatibility/scope, and PR-readability reviews before publication.

@vercel

vercel Bot commented Jul 19, 2026

Copy link
Copy Markdown

@wjueyao is attempting to deploy a commit to the IndexLabs Team on Vercel.

A member of the Team first needs to authorize it.

@multica-eve multica-eve changed the title feat(daemon): add drain-aware restart for active tasks MUL-4976: feat(daemon): add drain-aware restart for active tasks Jul 20, 2026

@multica-eve multica-eve left a comment

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

此前,管理员在已有工作进行时更新本机运行程序,只能立即打断这些工作;改动后,管理员可以选择先停止接收新工作,等已开始的工作自然结束后再更新。原来的立即更新方式保持不变。

产品上应该接受。 这个可选能力能避免维护操作制造失败尝试、重复执行和额外模型消耗。

🔄 请求修改:等待期间仍可能被服务端下发的更新打断。

server/internal/daemon/daemon.go:2699 的 drain barrier 只与周期性 auto-update 协调;心跳触发的 handleUpdate:2529)不检查或持有该 barrier,并在 :2574 直接调用 triggerRestart()。因此两个并发顺序都会破坏本功能的核心保证:

  • drain 已开始后收到 pending update,会取消仍在运行的任务;
  • pending update 先开始时,drain 仍会成功进入,随后被更新重启打断。

请将手动 drain、周期性 auto-update、心跳触发 update 纳入同一个带 owner 的互斥协议:drain 持有时 update 应延后;update 已开始时 drain 应返回冲突。请补齐 drain → updateupdate → drain 两个顺序的并发测试,证明任何一方都不能绕过另一方的所有权。

wjueyao added 2 commits July 20, 2026 17:23
Use one owner-checked claim barrier for manual drain, periodic auto-update, and heartbeat-triggered update. Cover both concurrency orderings with deterministic regression tests.
@wjueyao

wjueyao commented Jul 20, 2026

Copy link
Copy Markdown
Contributor Author

Addressed the drain/update concurrency gap in 8e4e188b5.

  • Manual drain, periodic auto-update, and heartbeat-triggered update now use one owner-checked claim barrier under claimMu.
  • A heartbeat update defers without reporting failure while drain owns the barrier, so it remains pending for a later heartbeat.
  • A drain request now returns 409 Conflict while an update owns the barrier.
  • Failure/cancellation paths release only their own owner; a scheduled restart retains ownership through process shutdown.
  • Added deterministic regression coverage for both drain -> heartbeat update and heartbeat update -> drain orderings, plus mismatched-owner release coverage.

Local verification:

go test ./internal/daemon ./cmd/multica -count=1
go test -race ./internal/daemon ./cmd/multica -run 'Test(ShutdownHandler|TryEnterClaim_RespectsBarrier|TryAutoUpdate|HandleUpdateDefersWhileDrainOwnsBarrier|DrainReturnsConflictWhileHandleUpdateOwnsBarrier|HealthHandlerReportsCLIVersionAndActiveTaskCount|DaemonRestartCommandExposesDrainFlag|RequestDaemonDrainShutdown)' -count=1
go vet ./internal/daemon ./cmd/multica
gofmt -l ...  # no output

The new CI run is in progress.

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

[Feature]: allow daemon restart to drain active tasks without cancelling them

2 participants