Skip to content

Commit b62bb8d

Browse files
committed
fix(security): propagate parent context to delegate_tasks so sub-agents die on cancel (M3)
delegateTasksTool.runTask() created its per-task context from context.Background(), completely detached from the parent agent. When the parent was cancelled (Ctrl+C, restart, timeout), the sub-agent processes kept running for the full 120s timeout, blocking the parent loop. Fix: - Add ctx field + SetContext() to delegateTasksTool - runTask now derives its timeout context from the parent's context instead of context.Background() - In the agent loop (loop.go:741), before calling t.Call(), check if the tool implements SetContext(context.Context) and propagate the agent's context Result: when the parent is cancelled, all running sub-agents are killed promptly via exec.CommandContext cancellation chain, instead of running the full 120s timeout. Fixes M3 from security audit.
1 parent 6c852ab commit b62bb8d

2 files changed

Lines changed: 26 additions & 1 deletion

File tree

cmd/odek/subagent_tool.go

Lines changed: 21 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -26,6 +26,12 @@ type delegateTasksTool struct {
2626
odekPath string // path to the odek binary
2727
timeout time.Duration
2828

29+
// ctx is the parent agent's context, set by the agent loop before each
30+
// Call invocation. When the parent is cancelled (Ctrl+C, restart, timeout),
31+
// runTask derives its per-task context from this, so sub-agent processes
32+
// are killed promptly instead of running the full timeout.
33+
ctx context.Context
34+
2935
// OnSubagentLog, if set, is called with each NDJSON progress line
3036
// emitted by a sub-agent. taskIdx is the index within the current
3137
// batch. Used by the WebUI for live log streaming.
@@ -34,6 +40,13 @@ type delegateTasksTool struct {
3440

3541
func (t *delegateTasksTool) Name() string { return "delegate_tasks" }
3642

43+
// SetContext sets the parent agent's context on the tool.
44+
// Called by the agent loop before each Call invocation to propagate
45+
// cancellation signals (Ctrl+C, restart, timeout) to sub-agents.
46+
func (t *delegateTasksTool) SetContext(ctx context.Context) {
47+
t.ctx = ctx
48+
}
49+
3750
func (t *delegateTasksTool) Description() string {
3851
return `Spawn one or more sub-agent OS processes to work on focused sub-tasks in parallel. Each sub-agent gets its own process, config, and context window. Use this when the task has clear independent sub-tasks that can be worked on simultaneously.
3952
@@ -142,7 +155,14 @@ func (t *delegateTasksTool) Call(args string) (string, error) {
142155
}
143156

144157
func (t *delegateTasksTool) runTask(taskIdx int, goal, taskContext, system string) string {
145-
ctx, cancel := context.WithTimeout(context.Background(), t.timeout)
158+
// Derive per-task context from the parent's context (if set).
159+
// When the parent is cancelled, all running sub-agents are killed
160+
// promptly instead of running the full timeout.
161+
parentCtx := context.Background()
162+
if t.ctx != nil {
163+
parentCtx = t.ctx
164+
}
165+
ctx, cancel := context.WithTimeout(parentCtx, t.timeout)
146166
defer cancel()
147167

148168
// Write task to temp file (avoids CLI arg length limits)

internal/loop/loop.go

Lines changed: 5 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -738,6 +738,11 @@ if e.approver != nil && len(result.ToolCalls) > 1 {
738738
t := e.registry.Get(tcRef.Function.Name)
739739
output := fmt.Sprintf("error: tool %q not found", tcRef.Function.Name)
740740
if t != nil {
741+
// Propagate agent context to tools that support it
742+
// (e.g. delegate_tasks kills sub-agents on parent cancel).
743+
if ctxTool, ok := t.(interface{ SetContext(context.Context) }); ok {
744+
ctxTool.SetContext(ctx)
745+
}
741746
res, err := t.Call(tcRef.Function.Arguments)
742747
if err != nil {
743748
output = fmt.Sprintf("error: %s", err.Error())

0 commit comments

Comments
 (0)