Add files via upload

internal/c2/hitl_context.go

@@ -0,0 +1,29 @@
+package c2
+
+import "context"
+
+type hitlRunCtxKey struct{}
+
+// WithHITLRunContext 将 runCtx（通常为整条 Agent / SSE 请求生命周期）挂到传入的 ctx 上。
+// MCP 工具 handler 收到的 ctx 可能是带单次工具超时的子 context，在工具 return 时会被 cancel；
+// 危险任务 HITL 应通过 HITLUserContext 使用 runCtx 等待人工审批。
+func WithHITLRunContext(ctx, runCtx context.Context) context.Context {
+	if ctx == nil || runCtx == nil {
+		return ctx
+	}
+	return context.WithValue(ctx, hitlRunCtxKey{}, runCtx)
+}
+
+// HITLUserContext 返回用于 C2 危险任务 HITL 等待的 context：
+// 若曾用 WithHITLRunContext 注入更长寿命的 runCtx 则返回之，否则返回 ctx。
+func HITLUserContext(ctx context.Context) context.Context {
+	if ctx == nil {
+		return context.Background()
+	}
+	if v := ctx.Value(hitlRunCtxKey{}); v != nil {
+		if run, ok := v.(context.Context); ok && run != nil {
+			return run
+		}
+	}
+	return ctx
+}